feat: update managed-agent to use track_metrics_of_async and add provider helper tests

feat: add TestGetAIUsageFromResponse and TestGetToolCallsFromResponse test coverage for LangChainHelper
feat: add TestGetAIUsageFromResponse test coverage for OpenAIHelper
fix: update ManagedAgent.invoke to use track_metrics_of_async

Author: jsonbailey

packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py

@@ -127,6 +127,71 @@ def test_creates_metrics_with_success_true_and_no_usage_when_metadata_missing(se
         assert result.success is True
         assert result.usage is None
 
+    def test_usage_metadata_preferred_over_response_metadata(self):
+        """usage_metadata should be used when it has non-zero counts."""
+        mock_response = AIMessage(content='Test')
+        mock_response.usage_metadata = {
+            'total_tokens': 10,
+            'input_tokens': 4,
+            'output_tokens': 6,
+        }
+        mock_response.response_metadata = {
+            'tokenUsage': {
+                'totalTokens': 999,
+                'promptTokens': 500,
+                'completionTokens': 499,
+            },
+        }
+        usage = LangChainHelper.get_ai_usage_from_response(mock_response)
+        assert usage is not None
+        assert usage.total == 10
+        assert usage.input == 4
+        assert usage.output == 6
+
+
+class TestGetAIUsageFromResponse:
+    """Tests for LangChainHelper.get_ai_usage_from_response."""
+
+    def test_returns_none_when_no_usage(self):
+        msg = AIMessage(content='hi')
+        assert LangChainHelper.get_ai_usage_from_response(msg) is None
+
+    def test_returns_none_when_all_zeros_in_metadata(self):
+        msg = AIMessage(content='hi')
+        msg.usage_metadata = {'total_tokens': 0, 'input_tokens': 0, 'output_tokens': 0}
+        assert LangChainHelper.get_ai_usage_from_response(msg) is None
+
+
+class TestGetToolCallsFromResponse:
+    """Tests for LangChainHelper.get_tool_calls_from_response."""
+
+    def test_returns_empty_when_no_tool_calls(self):
+        msg = AIMessage(content='hi')
+        assert LangChainHelper.get_tool_calls_from_response(msg) == []
+
+    def test_returns_empty_when_tool_calls_not_a_sequence(self):
+        msg = AIMessage(content='hi')
+        msg.tool_calls = None  # type: ignore
+        assert LangChainHelper.get_tool_calls_from_response(msg) == []
+
+    def test_extracts_names_from_dict_tool_calls(self):
+        msg = AIMessage(content='')
+        msg.tool_calls = [  # type: ignore
+            {'name': 'search', 'args': {}, 'id': '1'},
+            {'name': 'calc', 'args': {}, 'id': '2'},
+        ]
+        assert LangChainHelper.get_tool_calls_from_response(msg) == ['search', 'calc']
+
+    def test_returns_empty_when_tool_calls_is_not_a_list(self):
+        msg = AIMessage(content='hi')
+        msg.tool_calls = ()  # type: ignore
+        assert LangChainHelper.get_tool_calls_from_response(msg) == []
+
+    def test_skips_entries_without_name(self):
+        msg = AIMessage(content='')
+        msg.tool_calls = [{'name': 'a', 'id': '1'}, {}, {'name': 'b', 'id': '2'}]  # type: ignore
+        assert LangChainHelper.get_tool_calls_from_response(msg) == ['a', 'b']
+
 
 class TestMapProvider:
     """Tests for map_provider."""

packages/ai-providers/server-ai-openai/tests/test_openai_provider.py

@@ -8,6 +8,35 @@
 from ldai_openai import OpenAIModelRunner, OpenAIRunnerFactory, get_ai_metrics_from_response
 
 
+class TestGetAIUsageFromResponse:
+    """Tests for OpenAIHelper.get_ai_usage_from_response."""
+
+    def test_returns_usage_when_present(self):
+        mock_response = MagicMock()
+        mock_response.usage = MagicMock()
+        mock_response.usage.prompt_tokens = 50
+        mock_response.usage.completion_tokens = 50
+        mock_response.usage.total_tokens = 100
+        u = OpenAIHelper.get_ai_usage_from_response(mock_response)
+        assert u is not None
+        assert u.total == 100
+        assert u.input == 50
+        assert u.output == 50
+
+    def test_returns_none_when_usage_missing(self):
+        mock_response = MagicMock()
+        mock_response.usage = None
+        assert OpenAIHelper.get_ai_usage_from_response(mock_response) is None
+
+    def test_returns_none_when_all_counts_zero(self):
+        mock_response = MagicMock()
+        mock_response.usage = MagicMock()
+        mock_response.usage.total_tokens = 0
+        mock_response.usage.prompt_tokens = 0
+        mock_response.usage.completion_tokens = 0
+        assert OpenAIHelper.get_ai_usage_from_response(mock_response) is None
+
+
 class TestGetAIMetricsFromResponse:
     """Tests for get_ai_metrics_from_response."""
 

packages/sdk/server-ai/src/ldai/managed_agent.py

@@ -31,7 +31,7 @@ async def run(self, input: str) -> AgentResult:
         :param input: The user prompt or input to the agent
         :return: AgentResult containing the agent's output and metrics
         """
-        return await self._tracker.track_metrics_of(
+        return await self._tracker.track_metrics_of_async(
             lambda: self._agent_runner.run(input),
             lambda result: result.metrics,
         )