unique_orchestrator 2026.30.0.dev7__tar.gz → 2026.30.0.dev8__tar.gz

{unique_orchestrator-2026.30.0.dev7 → unique_orchestrator-2026.30.0.dev8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-orchestrator
-Version: 2026.30.0.dev7
+Version: 2026.30.0.dev8
 Summary: 
 Author: Andreas Hauri
 Author-email: Andreas Hauri <andreas.hauri@unique.ai>
@@ -9,9 +9,9 @@ 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.30.0.dev19,<2026.30.0rc0
+Requires-Dist: unique-toolkit>=2026.30.0.dev20,<2026.30.0rc0
 Requires-Dist: unique-stock-ticker>=2026.30.0.dev0,<2026.30.0rc0
-Requires-Dist: unique-follow-up-questions>=2026.30.0.dev0,<2026.30.0rc0
+Requires-Dist: unique-follow-up-questions>=2026.30.0.dev1,<2026.30.0rc0
 Requires-Dist: unique-internal-search>=2026.30.0.dev0,<2026.30.0rc0
 Requires-Dist: unique-deep-research>=2026.30.0.dev0,<2026.30.0rc0
 Requires-Dist: unique-web-search>=2026.30.0.dev3,<2026.30.0rc0

{unique_orchestrator-2026.30.0.dev7 → unique_orchestrator-2026.30.0.dev8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_orchestrator"
-version = "2026.30.0.dev7"
+version = "2026.30.0.dev8"
 description = ""
 readme = "README.md"
 license = { text = "Proprietary" }
@@ -13,9 +13,9 @@ dependencies = [
     "pydantic-settings>=2.10.1,<3",
     "typing-extensions>=4.9.0,<5",
     "jinja2>=3.1.0,<4",
-    "unique-toolkit>=2026.30.0.dev19,<2026.30.0rc0",
+    "unique-toolkit>=2026.30.0.dev20,<2026.30.0rc0",
     "unique-stock-ticker>=2026.30.0.dev0,<2026.30.0rc0",
-    "unique-follow-up-questions>=2026.30.0.dev0,<2026.30.0rc0",
+    "unique-follow-up-questions>=2026.30.0.dev1,<2026.30.0rc0",
     "unique-internal-search>=2026.30.0.dev0,<2026.30.0rc0",
     "unique-deep-research>=2026.30.0.dev0,<2026.30.0rc0",
     "unique-web-search>=2026.30.0.dev3,<2026.30.0rc0",

{unique_orchestrator-2026.30.0.dev7 → unique_orchestrator-2026.30.0.dev8}/unique_orchestrator/tests/test_unique_ai_async_modify.py

@@ -56,10 +56,12 @@ def _make_ua(monkeypatch, *, feature_flag_enabled: bool = False):
     mock_postprocessor_manager = MagicMock()
     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 = []
 
     mock_evaluation_manager = MagicMock()
     mock_evaluation_manager.run_evaluations = AsyncMock(return_value=[])
     mock_evaluation_manager.get_execution_times.return_value = {}
+    mock_evaluation_manager.get_invocation_stats.return_value = []
 
     dummy_event = MagicMock()
     dummy_event.payload.assistant_message.id = "assist_1"
@@ -89,8 +91,10 @@ def _make_ua(monkeypatch, *, feature_flag_enabled: bool = False):
     ua._execution_times = []
     ua._current_loop_timing = {}
     ua._generated_files_info = None
+    ua._invocation_stats = []
     ua.current_iteration_index = 0
     ua._skill_choices = []
+    ua._loop_iteration_runner = MagicMock()
 
     return ua
 
@@ -107,6 +111,7 @@ async def test_run_calls_modify_async_when_feature_flag_disabled(monkeypatch):
     empty_response.message.references = []
     empty_response.tool_calls = None
     empty_response.is_empty.return_value = True
+    empty_response.usage = None
 
     ua._plan_or_execute = AsyncMock(return_value=empty_response)
 

{unique_orchestrator-2026.30.0.dev7 → unique_orchestrator-2026.30.0.dev8}/unique_orchestrator/tests/test_unique_ai_execution_timing.py

@@ -65,6 +65,7 @@ def _make_tool_response(name: str, execution_time_s: float | None = None) -> Mag
     response.debug_info = (
         {"execution_time_s": execution_time_s} if execution_time_s is not None else None
     )
+    response.invocation_stats = []
     return response
 
 
@@ -75,6 +76,7 @@ def _make_loop_response(tool_calls: list | None = None) -> MagicMock:
     response.message.references = []
     response.message.text = "response text"
     response.is_empty.return_value = False
+    response.usage = None
     return response
 
 
@@ -397,9 +399,11 @@ class TestHandleNoToolCallsTiming:
             "evaluation": {},
         }
         instance._evaluation_manager.run_evaluations = AsyncMock(return_value=[])
+        instance._evaluation_manager.get_invocation_stats.return_value = []
         instance._postprocessor_manager.run_postprocessors = AsyncMock(
             return_value=None
         )
+        instance._postprocessor_manager.get_invocation_stats.return_value = []
         return instance
 
     @pytest.mark.ai
@@ -666,10 +670,12 @@ class TestRunExecutionTimingIntegration:
         mock_postprocessor_manager = MagicMock()
         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 = []
 
         mock_evaluation_manager = MagicMock()
         mock_evaluation_manager.run_evaluations = AsyncMock(return_value=[])
         mock_evaluation_manager.get_execution_times.return_value = {}
+        mock_evaluation_manager.get_invocation_stats.return_value = []
 
         dummy_event = MagicMock()
         dummy_event.payload.assistant_message.id = "assist_1"

{unique_orchestrator-2026.30.0.dev7 → unique_orchestrator-2026.30.0.dev8}/unique_orchestrator/tests/test_unique_ai_loop_debug_params.py

@@ -78,6 +78,7 @@ def _make_loop_response() -> MagicMock:
     response.message.text = "response text"
     response.message.original_text = "original response text"
     response.is_empty.return_value = False
+    response.usage = None
     return response
 
 

{unique_orchestrator-2026.30.0.dev7 → unique_orchestrator-2026.30.0.dev8}/unique_orchestrator/tests/test_unique_ai_reference_order.py

@@ -68,6 +68,7 @@ async def test_history_updated_before_reference_extraction(monkeypatch):
             self.message = MagicMock()
             self.message.references = []
             self.message.text = ""
+            self.usage = None
 
         def is_empty(self):
             return False

unique_orchestrator-2026.30.0.dev8/unique_orchestrator/tests/test_unique_ai_token_usage.py

@@ -0,0 +1,480 @@
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+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.schemas import LanguageModelTokenUsage
+
+from unique_orchestrator.unique_ai import UniqueAI
+
+MODEL_NAME = "gpt-4"
+
+
+def _make_loop_response(usage: LanguageModelTokenUsage | None) -> MagicMock:
+    response = MagicMock()
+    response.tool_calls = None
+    response.message = MagicMock()
+    response.message.references = []
+    response.message.text = "response text"
+    response.is_empty.return_value = False
+    response.usage = usage
+    return response
+
+
+def _build_run_ua(loop_responses: list[MagicMock], max_iterations: int):
+    mock_cancellation = MagicMock()
+    mock_cancellation.is_cancelled = False
+    mock_cancellation.on_cancellation.subscribe = MagicMock(return_value=MagicMock())
+    mock_cancellation.check_cancellation_async = AsyncMock(return_value=False)
+
+    mock_chat_service = MagicMock()
+    mock_chat_service.cancellation = mock_cancellation
+    mock_chat_service.get_debug_info_async = AsyncMock(return_value={})
+    mock_chat_service.update_debug_info_async = AsyncMock(return_value=None)
+    mock_chat_service.modify_assistant_message_async = AsyncMock(return_value=None)
+    mock_chat_service.create_assistant_message_async = AsyncMock(
+        return_value=MagicMock(id="assist_new")
+    )
+
+    mock_tool_manager = MagicMock()
+    mock_tool_manager.available_tools = []
+    mock_tool_manager.get_forced_tools.return_value = []
+    mock_tool_manager.get_tool_definitions.return_value = []
+    mock_tool_manager.filter_duplicate_tool_calls.side_effect = lambda x: x
+    mock_tool_manager.filter_tool_calls_by_max_tool_calls_allowed.side_effect = (
+        lambda x: x
+    )
+    mock_tool_manager.does_a_tool_take_control.return_value = False
+    mock_tool_manager.should_stop_after_tool_calls.return_value = False
+    mock_tool_manager.get_evaluation_check_list.return_value = []
+
+    mock_debug_info_manager = MagicMock()
+    mock_debug_info_manager.get.return_value = {"tools": []}
+
+    mock_config = MagicMock()
+    mock_config.effective_max_loop_iterations = max_iterations
+    mock_config.agent.prompt_config.user_metadata = []
+    mock_config.space.language_model.name = MODEL_NAME
+
+    mock_history_manager = MagicMock()
+    mock_history_manager.get_history_for_model_call = AsyncMock(
+        return_value=MagicMock()
+    )
+    mock_history_manager._append_tool_calls_to_history = MagicMock()
+    mock_history_manager.add_tool_call_results = MagicMock()
+    mock_history_manager.extract_message_tools.return_value = []
+
+    mock_postprocessor_manager = MagicMock()
+    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 = []
+
+    mock_evaluation_manager = MagicMock()
+    mock_evaluation_manager.run_evaluations = AsyncMock(return_value=[])
+    mock_evaluation_manager.get_execution_times.return_value = {}
+    mock_evaluation_manager.get_invocation_stats.return_value = []
+
+    dummy_event = MagicMock()
+    dummy_event.payload.assistant_message.id = "assist_1"
+    dummy_event.payload.user_message.text = "query"
+    dummy_event.payload.assistant_id = "assistant_123"
+    dummy_event.payload.name = "TestAssistant"
+    dummy_event.payload.user_metadata = {}
+    dummy_event.payload.tool_parameters = {}
+    dummy_event.company_id = "company_1"
+
+    ua = UniqueAI(
+        logger=MagicMock(),
+        event=dummy_event,
+        config=mock_config,
+        chat_service=mock_chat_service,
+        content_service=MagicMock(),
+        debug_info_manager=mock_debug_info_manager,
+        streaming_handler=MagicMock(),
+        reference_manager=MagicMock(),
+        thinking_manager=MagicMock(),
+        tool_manager=mock_tool_manager,
+        history_manager=mock_history_manager,
+        evaluation_manager=mock_evaluation_manager,
+        postprocessor_manager=mock_postprocessor_manager,
+        message_step_logger=MagicMock(),
+        mcp_servers=[],
+        loop_iteration_runner=AsyncMock(side_effect=loop_responses),
+    )
+    ua._render_user_prompt = AsyncMock(return_value="user")  # type: ignore[method-assign]
+    ua._render_system_prompt = AsyncMock(return_value="system")  # type: ignore[method-assign]
+    ua._thinking_manager.thinking_is_displayed = MagicMock(return_value=True)
+    ua._thinking_manager.update_start_text = MagicMock(return_value="")
+    return ua
+
+
+class TestAccumulateUsage:
+    """Tests for how UniqueAI._invocation_stats is populated from the main loop."""
+
+    def test_main_loop__none_usage__is_noop(self):
+        ua = _build_run_ua([_make_loop_response(None)], max_iterations=1)
+        assert ua._invocation_stats == []
+
+    @pytest.mark.asyncio
+    async def test_main_loop__single_call__appends_one_entry(self):
+        loop_response = _make_loop_response(
+            LanguageModelTokenUsage(
+                completion_tokens=10, prompt_tokens=20, total_tokens=30
+            )
+        )
+        ua = _build_run_ua([loop_response], max_iterations=1)
+
+        await ua.run()
+
+        assert len(ua._invocation_stats) == 1
+        entry = ua._invocation_stats[0]
+        assert entry.model_name == MODEL_NAME
+        assert entry.source == "main_loop[1]"
+        assert entry.token_usage == LanguageModelTokenUsage(
+            completion_tokens=10, prompt_tokens=20, total_tokens=30
+        )
+
+    @pytest.mark.asyncio
+    async def test_main_loop__multiple_iterations__appends_one_entry_per_call(self):
+        tool_call = MagicMock()
+        tool_call.name = "WebSearch"
+
+        first = _make_loop_response(
+            LanguageModelTokenUsage(
+                completion_tokens=10, prompt_tokens=20, total_tokens=30
+            )
+        )
+        first.tool_calls = [tool_call]
+        second = _make_loop_response(
+            LanguageModelTokenUsage(
+                completion_tokens=1, prompt_tokens=2, total_tokens=3
+            )
+        )
+        ua = _build_run_ua([first, second], max_iterations=2)
+        ua._tool_manager.execute_selected_tools = AsyncMock(
+            return_value=[
+                ToolCallResponse(id="call_1", name="WebSearch", invocation_stats=[])
+            ]
+        )
+
+        await ua.run()
+
+        assert len(ua._invocation_stats) == 2
+        assert [entry.token_usage.total_tokens for entry in ua._invocation_stats] == [
+            30,
+            3,
+        ]
+        assert [entry.source for entry in ua._invocation_stats] == [
+            "main_loop[1]",
+            "main_loop[2]",
+        ]
+
+    @pytest.mark.asyncio
+    async def test_main_loop__none_mixed_with_real__none_is_skipped(self):
+        tool_call = MagicMock()
+        tool_call.name = "WebSearch"
+
+        first = _make_loop_response(
+            LanguageModelTokenUsage(
+                completion_tokens=10, prompt_tokens=20, total_tokens=30
+            )
+        )
+        first.tool_calls = [tool_call]
+        second = _make_loop_response(None)
+        ua = _build_run_ua([first, second], max_iterations=2)
+        ua._tool_manager.execute_selected_tools = AsyncMock(
+            return_value=[
+                ToolCallResponse(id="call_1", name="WebSearch", invocation_stats=[])
+            ]
+        )
+
+        await ua.run()
+
+        assert len(ua._invocation_stats) == 1
+        assert ua._invocation_stats[0].token_usage == LanguageModelTokenUsage(
+            completion_tokens=10, prompt_tokens=20, total_tokens=30
+        )
+
+
+class TestRunTokenUsageIntegration:
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__debug_info_manager__receives_llm_invocations_key(self):
+        """A single-iteration run with real usage on the loop response must
+        surface it under debug_info's 'llm_invocations' key."""
+        loop_response = _make_loop_response(
+            LanguageModelTokenUsage(
+                completion_tokens=12, prompt_tokens=34, total_tokens=46
+            )
+        )
+        ua = _build_run_ua([loop_response], max_iterations=1)
+
+        await ua.run()
+
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        assert "llm_invocations" in calls_by_key
+        payload = calls_by_key["llm_invocations"]
+        assert len(payload) == 1
+        assert payload[0]["modelName"] == MODEL_NAME
+        assert payload[0]["source"] == "main_loop[1]"
+        assert payload[0]["tokenUsage"] == {
+            "completionTokens": 12,
+            "promptTokens": 34,
+            "totalTokens": 46,
+            "reasoningTokens": None,
+            "cachedTokens": None,
+            "cacheWriteTokens": None,
+        }
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__no_usage_on_any_response__llm_invocations_is_empty(
+        self,
+    ):
+        """When usage is never populated (e.g. no supporting provider),
+        the 'llm_invocations' key must still be present as an empty list,
+        not absent/silent."""
+        ua = _build_run_ua([_make_loop_response(None)], max_iterations=1)
+
+        await ua.run()
+
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        assert "llm_invocations" in calls_by_key
+        assert calls_by_key["llm_invocations"] == []
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__postprocessor_usage__added_to_loop_usage(self):
+        """FollowUpPostprocessor (and any postprocessor) makes its own LLM
+        call outside _plan_or_execute() -- its usage must appear as its own
+        entry in the list, not silently dropped."""
+        loop_response = _make_loop_response(
+            LanguageModelTokenUsage(
+                completion_tokens=12, prompt_tokens=34, total_tokens=46
+            )
+        )
+        ua = _build_run_ua([loop_response], max_iterations=1)
+        ua._postprocessor_manager.get_invocation_stats.return_value = [
+            LanguageModelInvocationStats.from_usage(
+                MODEL_NAME,
+                LanguageModelTokenUsage(
+                    completion_tokens=3, prompt_tokens=7, total_tokens=10
+                ),
+                source="FollowUpPostprocessor",
+            )
+        ]
+
+        await ua.run()
+
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        payload = calls_by_key["llm_invocations"]
+        sources = [inv["source"] for inv in payload]
+        assert sources == ["main_loop[1]", "FollowUpPostprocessor"]
+        assert payload[0]["tokenUsage"]["totalTokens"] == 46
+        assert payload[1]["tokenUsage"]["totalTokens"] == 10
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__loop_runner_reports_invocation_stats__added_to_loop_usage(
+        self,
+    ):
+        """PlanningMiddleware (and any loop-runner wrapper implementing
+        SupportsInvocationStats) makes its own LLM call outside the main
+        completion -- e.g. a planning step before the iteration -- and its
+        usage must appear as its own entry in the list, not silently dropped."""
+        loop_response = _make_loop_response(
+            LanguageModelTokenUsage(
+                completion_tokens=12, prompt_tokens=34, total_tokens=46
+            )
+        )
+        ua = _build_run_ua([loop_response], max_iterations=1)
+        # A bare AsyncMock() satisfies hasattr() for any name but does NOT
+        # satisfy isinstance() against a runtime_checkable Protocol -- only
+        # spec=<the protocol> does, matching what real PlanningMiddleware
+        # instances trigger via isinstance() in production.
+        runner = AsyncMock(spec=SupportsInvocationStats)
+        runner.side_effect = [loop_response]
+        runner.get_invocation_stats.return_value = [
+            LanguageModelInvocationStats.from_usage(
+                MODEL_NAME,
+                LanguageModelTokenUsage(
+                    completion_tokens=2, prompt_tokens=8, total_tokens=10
+                ),
+                source="planning",
+            )
+        ]
+        ua._loop_iteration_runner = runner
+
+        await ua.run()
+
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        payload = calls_by_key["llm_invocations"]
+        sources = [inv["source"] for inv in payload]
+        assert sources == ["main_loop[1]", "planning"]
+        assert payload[0]["tokenUsage"]["totalTokens"] == 46
+        assert payload[1]["tokenUsage"]["totalTokens"] == 10
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__loop_runner_without_invocation_stats__is_a_noop(self):
+        """A plain runner that doesn't implement SupportsInvocationStats
+        (the common case: Basic/Qwen/Mistral runners with no planning
+        middleware) must not contribute spurious entries."""
+        loop_response = _make_loop_response(
+            LanguageModelTokenUsage(
+                completion_tokens=12, prompt_tokens=34, total_tokens=46
+            )
+        )
+        ua = _build_run_ua([loop_response], max_iterations=1)
+        # spec restricts this mock to only __call__ -- no get_invocation_stats
+        # attribute exists, so isinstance(..., SupportsInvocationStats) is False,
+        # mirroring a real runner (Basic/Qwen/Mistral) that doesn't implement it.
+        ua._loop_iteration_runner = AsyncMock(
+            side_effect=[loop_response], spec=["__call__"]
+        )
+
+        await ua.run()
+
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        payload = calls_by_key["llm_invocations"]
+        sources = [inv["source"] for inv in payload]
+        assert sources == ["main_loop[1]"]
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__tool_usage__added_to_loop_usage(self):
+        """A tool's own internal LLM calls (e.g. web search argument screening)
+        report usage via ToolCallResponse.invocation_stats -- it must be folded
+        into the same total as the main loop, per tool-call iteration."""
+        tool_call = MagicMock()
+        tool_call.name = "WebSearch"
+
+        loop_response_with_tools = _make_loop_response(
+            LanguageModelTokenUsage(
+                completion_tokens=12, prompt_tokens=34, total_tokens=46
+            )
+        )
+        loop_response_with_tools.tool_calls = [tool_call]
+        loop_response_with_tools.is_empty.return_value = False
+
+        final_response = _make_loop_response(None)
+
+        ua = _build_run_ua([loop_response_with_tools, final_response], max_iterations=2)
+        ua._tool_manager.execute_selected_tools = AsyncMock(
+            return_value=[
+                ToolCallResponse(
+                    id="call_1",
+                    name="WebSearch",
+                    invocation_stats=[
+                        LanguageModelInvocationStats.from_usage(
+                            MODEL_NAME,
+                            LanguageModelTokenUsage(
+                                completion_tokens=1, prompt_tokens=2, total_tokens=3
+                            ),
+                            source="WebSearch",
+                        )
+                    ],
+                )
+            ]
+        )
+
+        await ua.run()
+
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        payload = calls_by_key["llm_invocations"]
+        sources = [inv["source"] for inv in payload]
+        assert sources == ["main_loop[1]", "WebSearch"]
+        assert payload[0]["tokenUsage"]["totalTokens"] == 46
+        assert payload[1]["tokenUsage"]["totalTokens"] == 3
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__multiple_tool_calls_in_one_iteration__usage_summed(self):
+        """Two (or more) tool calls in the same iteration -- e.g. two parallel
+        web searches -- must each contribute their own usage, not overwrite
+        or drop one another."""
+        tool_call = MagicMock()
+        tool_call.name = "WebSearch"
+
+        loop_response_with_tools = _make_loop_response(None)
+        loop_response_with_tools.tool_calls = [tool_call, tool_call]
+        loop_response_with_tools.is_empty.return_value = False
+
+        final_response = _make_loop_response(None)
+
+        ua = _build_run_ua([loop_response_with_tools, final_response], max_iterations=2)
+        ua._tool_manager.execute_selected_tools = AsyncMock(
+            return_value=[
+                ToolCallResponse(
+                    id="call_1",
+                    name="WebSearch",
+                    invocation_stats=[
+                        LanguageModelInvocationStats.from_usage(
+                            MODEL_NAME,
+                            LanguageModelTokenUsage(
+                                completion_tokens=1, prompt_tokens=2, total_tokens=3
+                            ),
+                            source="WebSearch",
+                        )
+                    ],
+                ),
+                ToolCallResponse(
+                    id="call_2",
+                    name="WebSearch",
+                    invocation_stats=[
+                        LanguageModelInvocationStats.from_usage(
+                            MODEL_NAME,
+                            LanguageModelTokenUsage(
+                                completion_tokens=4, prompt_tokens=5, total_tokens=9
+                            ),
+                            source="WebSearch",
+                        )
+                    ],
+                ),
+            ]
+        )
+
+        await ua.run()
+
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        payload = calls_by_key["llm_invocations"]
+        assert len(payload) == 2
+        assert payload[0]["tokenUsage"]["totalTokens"] == 3
+        assert payload[1]["tokenUsage"]["totalTokens"] == 9
+
+    @pytest.mark.ai
+    @pytest.mark.asyncio
+    async def test_run__reset_between_runs__does_not_leak_across_calls(self):
+        """_invocation_stats must reset at the start of run(), so a second run()
+        on the same UniqueAI instance doesn't inherit the prior run's entries."""
+        loop_response = _make_loop_response(
+            LanguageModelTokenUsage(
+                completion_tokens=5, prompt_tokens=5, total_tokens=10
+            )
+        )
+        ua = _build_run_ua([loop_response, loop_response], max_iterations=1)
+
+        await ua.run()
+        await ua.run()
+
+        calls_by_key = {
+            args[0][0]: args[0][1] for args in ua._debug_info_manager.add.call_args_list
+        }
+        payload = calls_by_key["llm_invocations"]
+        assert len(payload) == 1
+        assert payload[0]["tokenUsage"]["totalTokens"] == 10

{unique_orchestrator-2026.30.0.dev7 → unique_orchestrator-2026.30.0.dev8}/unique_orchestrator/unique_ai.py

@@ -21,6 +21,7 @@ from unique_toolkit.agentic.history_manager.utils import (
 from unique_toolkit.agentic.loop_runner import (
     LoopIterationRunner,
     ResponsesLoopIterationRunner,
+    SupportsInvocationStats,
 )
 from unique_toolkit.agentic.message_log_manager.service import MessageStepLogger
 from unique_toolkit.agentic.postprocessor.postprocessor_manager import (
@@ -49,6 +50,9 @@ from unique_toolkit.chat.service import ChatService
 from unique_toolkit.content import Content
 from unique_toolkit.content.service import ContentService
 from unique_toolkit.language_model import LanguageModelAssistantMessage
+from unique_toolkit.language_model.invocation_stats import (
+    LanguageModelInvocationStats,
+)
 from unique_toolkit.language_model.schemas import (
     LanguageModelMessages,
     LanguageModelStreamResponse,
@@ -204,6 +208,7 @@ class UniqueAI:
         self._current_loop_timing: dict[str, Any] = {}
         self._loop_debug_params: list[dict[str, Any]] = []
         self._generated_files_info: ArtifactsDebugInfo | None = None
+        self._invocation_stats: list[LanguageModelInvocationStats] = []
 
     async def _on_cancellation(self, _event: CancellationEvent) -> None:
         """Subscriber called by the cancellation event bus."""
@@ -233,6 +238,7 @@ class UniqueAI:
         # reaching _handle_no_tool_calls (tool takes control / empty response /
         # cancellation) must not report the previous run's artifacts.
         self._generated_files_info = None
+        self._invocation_stats = []
         run_start = time.perf_counter()
 
         await preload_invoked_skills(
@@ -268,6 +274,27 @@ class UniqueAI:
                     time.perf_counter() - planning_start, 3
                 )
                 self._logger.info("Done with _plan_or_execute")
+                if loop_response.usage is not None:
+                    self._invocation_stats.append(
+                        LanguageModelInvocationStats.from_usage(
+                            self._config.space.language_model.name,
+                            loop_response.usage,
+                            source=f"main_loop[{i + 1}]",
+                        )
+                    )
+                # TODO(UN-20907): if _plan_or_execute() above raises an exception that
+                # propagates out of this loop uncaught, execution never reaches this
+                # drain call -- any planning-step usage recorded on the runner this
+                # iteration is lost. In practice this is not planning-specific: an
+                # uncaught abort here also skips the debug_info["llm_invocations"]
+                # write entirely (see the `finally` at the end of run()), so every
+                # usage source (main_loop, evaluations, postprocessors) is equally
+                # absent from that turn, not just planning. Revisit only if partial
+                # debug_info persistence on an aborted run becomes a requirement.
+                if isinstance(self._loop_iteration_runner, SupportsInvocationStats):
+                    self._invocation_stats.extend(
+                        self._loop_iteration_runner.get_invocation_stats()
+                    )
 
                 if await self._chat_service.cancellation.check_cancellation_async():
                     self._finalize_loop_timing(loop_start)
@@ -329,6 +356,13 @@ 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)
+            self._debug_info_manager.add(
+                "llm_invocations",
+                [
+                    invocation.model_dump(by_alias=True)
+                    for invocation in self._invocation_stats
+                ],
+            )
             reference_sources = {
                 ("url", reference.url)
                 if reference.url is not None
@@ -724,6 +758,9 @@ class UniqueAI:
         self._current_loop_timing["post_processing"].update(
             self._postprocessor_manager.get_execution_times()
         )
+        self._invocation_stats.extend(
+            self._postprocessor_manager.get_invocation_stats()
+        )
 
         evaluation_times = self._evaluation_manager.get_execution_times()
         for name in selected_evaluation_names:
@@ -731,6 +768,7 @@ class UniqueAI:
             self._current_loop_timing["evaluation"][name_str] = evaluation_times.get(
                 name_str, 0
             )
+        self._invocation_stats.extend(self._evaluation_manager.get_invocation_stats())
 
         if evaluation_results.success and not all(
             result.is_positive for result in evaluation_results.unpack()
@@ -838,6 +876,9 @@ class UniqueAI:
         )
         execution_total = round(time.perf_counter() - execution_start, 3)
 
+        for response in tool_call_responses:
+            self._invocation_stats.extend(response.invocation_stats)
+
         tool_times: dict[str, float] = {}
         for response in tool_call_responses:
             if response.debug_info and "execution_time_s" in response.debug_info: