feat: extract mcp tool usage (#31802)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/tests/unit_tests/tools/test_mcp_tool.py

@@ -1,4 +1,5 @@
 import base64
+from decimal import Decimal
 from unittest.mock import Mock, patch
 
 import pytest
@@ -9,8 +10,10 @@ from core.mcp.types import (
     CallToolResult,
     EmbeddedResource,
     ImageContent,
+    TextContent,
     TextResourceContents,
 )
+from core.model_runtime.entities.llm_entities import LLMUsage
 from core.tools.__base.tool_runtime import ToolRuntime
 from core.tools.entities.common_entities import I18nObject
 from core.tools.entities.tool_entities import ToolEntity, ToolIdentity, ToolInvokeMessage
@@ -120,3 +123,231 @@ class TestMCPToolInvoke:
         # Validate values
         values = {m.message.variable_name: m.message.variable_value for m in var_msgs}
         assert values == {"a": 1, "b": "x"}
+
+
+class TestMCPToolUsageExtraction:
+    """Test usage metadata extraction from MCP tool results."""
+
+    def test_extract_usage_dict_from_direct_usage_field(self) -> None:
+        """Test extraction when usage is directly in meta.usage field."""
+        meta = {
+            "usage": {
+                "prompt_tokens": 100,
+                "completion_tokens": 50,
+                "total_tokens": 150,
+                "total_price": "0.001",
+                "currency": "USD",
+            }
+        }
+        usage_dict = MCPTool._extract_usage_dict(meta)
+        assert usage_dict is not None
+        assert usage_dict["prompt_tokens"] == 100
+        assert usage_dict["completion_tokens"] == 50
+        assert usage_dict["total_tokens"] == 150
+        assert usage_dict["total_price"] == "0.001"
+        assert usage_dict["currency"] == "USD"
+
+    def test_extract_usage_dict_from_nested_metadata(self) -> None:
+        """Test extraction when usage is nested in meta.metadata.usage."""
+        meta = {
+            "metadata": {
+                "usage": {
+                    "prompt_tokens": 200,
+                    "completion_tokens": 100,
+                    "total_tokens": 300,
+                }
+            }
+        }
+        usage_dict = MCPTool._extract_usage_dict(meta)
+        assert usage_dict is not None
+        assert usage_dict["prompt_tokens"] == 200
+        assert usage_dict["total_tokens"] == 300
+
+    def test_extract_usage_dict_from_flat_token_fields(self) -> None:
+        """Test extraction when token counts are directly in meta."""
+        meta = {
+            "prompt_tokens": 150,
+            "completion_tokens": 75,
+            "total_tokens": 225,
+            "currency": "EUR",
+        }
+        usage_dict = MCPTool._extract_usage_dict(meta)
+        assert usage_dict is not None
+        assert usage_dict["prompt_tokens"] == 150
+        assert usage_dict["completion_tokens"] == 75
+        assert usage_dict["total_tokens"] == 225
+        assert usage_dict["currency"] == "EUR"
+
+    def test_extract_usage_dict_recursive(self) -> None:
+        """Test recursive search through nested structures."""
+        meta = {
+            "custom": {
+                "nested": {
+                    "usage": {
+                        "total_tokens": 500,
+                        "prompt_tokens": 300,
+                        "completion_tokens": 200,
+                    }
+                }
+            }
+        }
+        usage_dict = MCPTool._extract_usage_dict(meta)
+        assert usage_dict is not None
+        assert usage_dict["total_tokens"] == 500
+
+    def test_extract_usage_dict_from_list(self) -> None:
+        """Test extraction from nested list structures."""
+        meta = {
+            "items": [
+                {"usage": {"total_tokens": 100}},
+                {"other": "data"},
+            ]
+        }
+        usage_dict = MCPTool._extract_usage_dict(meta)
+        assert usage_dict is not None
+        assert usage_dict["total_tokens"] == 100
+
+    def test_extract_usage_dict_returns_none_when_missing(self) -> None:
+        """Test that None is returned when no usage data is present."""
+        meta = {"other": "data", "custom": {"nested": {"value": 123}}}
+        usage_dict = MCPTool._extract_usage_dict(meta)
+        assert usage_dict is None
+
+    def test_extract_usage_dict_empty_meta(self) -> None:
+        """Test with empty meta dict."""
+        usage_dict = MCPTool._extract_usage_dict({})
+        assert usage_dict is None
+
+    def test_derive_usage_from_result_with_meta(self) -> None:
+        """Test _derive_usage_from_result with populated meta."""
+        meta = {
+            "usage": {
+                "prompt_tokens": 100,
+                "completion_tokens": 50,
+                "total_tokens": 150,
+                "total_price": "0.0015",
+                "currency": "USD",
+            }
+        }
+        result = CallToolResult(content=[], _meta=meta)
+        usage = MCPTool._derive_usage_from_result(result)
+
+        assert isinstance(usage, LLMUsage)
+        assert usage.prompt_tokens == 100
+        assert usage.completion_tokens == 50
+        assert usage.total_tokens == 150
+        assert usage.total_price == Decimal("0.0015")
+        assert usage.currency == "USD"
+
+    def test_derive_usage_from_result_without_meta(self) -> None:
+        """Test _derive_usage_from_result with no meta returns empty usage."""
+        result = CallToolResult(content=[], meta=None)
+        usage = MCPTool._derive_usage_from_result(result)
+
+        assert isinstance(usage, LLMUsage)
+        assert usage.total_tokens == 0
+        assert usage.prompt_tokens == 0
+        assert usage.completion_tokens == 0
+
+    def test_derive_usage_from_result_calculates_total_tokens(self) -> None:
+        """Test that total_tokens is calculated when missing."""
+        meta = {
+            "usage": {
+                "prompt_tokens": 100,
+                "completion_tokens": 50,
+                # total_tokens is missing
+            }
+        }
+        result = CallToolResult(content=[], _meta=meta)
+        usage = MCPTool._derive_usage_from_result(result)
+
+        assert usage.total_tokens == 150  # 100 + 50
+        assert usage.prompt_tokens == 100
+        assert usage.completion_tokens == 50
+
+    def test_invoke_sets_latest_usage_from_meta(self) -> None:
+        """Test that _invoke sets _latest_usage from result meta."""
+        tool = _make_mcp_tool()
+        meta = {
+            "usage": {
+                "prompt_tokens": 200,
+                "completion_tokens": 100,
+                "total_tokens": 300,
+                "total_price": "0.003",
+                "currency": "USD",
+            }
+        }
+        result = CallToolResult(content=[TextContent(type="text", text="test")], _meta=meta)
+
+        with patch.object(tool, "invoke_remote_mcp_tool", return_value=result):
+            list(tool._invoke(user_id="test_user", tool_parameters={}))
+
+        # Verify latest_usage was set correctly
+        assert tool.latest_usage.prompt_tokens == 200
+        assert tool.latest_usage.completion_tokens == 100
+        assert tool.latest_usage.total_tokens == 300
+        assert tool.latest_usage.total_price == Decimal("0.003")
+
+    def test_invoke_with_no_meta_returns_empty_usage(self) -> None:
+        """Test that _invoke returns empty usage when no meta is present."""
+        tool = _make_mcp_tool()
+        result = CallToolResult(content=[TextContent(type="text", text="test")], _meta=None)
+
+        with patch.object(tool, "invoke_remote_mcp_tool", return_value=result):
+            list(tool._invoke(user_id="test_user", tool_parameters={}))
+
+        # Verify latest_usage is empty
+        assert tool.latest_usage.total_tokens == 0
+        assert tool.latest_usage.prompt_tokens == 0
+        assert tool.latest_usage.completion_tokens == 0
+
+    def test_latest_usage_property_returns_llm_usage(self) -> None:
+        """Test that latest_usage property returns LLMUsage instance."""
+        tool = _make_mcp_tool()
+        assert isinstance(tool.latest_usage, LLMUsage)
+
+    def test_initial_usage_is_empty(self) -> None:
+        """Test that MCPTool is initialized with empty usage."""
+        tool = _make_mcp_tool()
+        assert tool.latest_usage.total_tokens == 0
+        assert tool.latest_usage.prompt_tokens == 0
+        assert tool.latest_usage.completion_tokens == 0
+        assert tool.latest_usage.total_price == Decimal(0)
+
+    @pytest.mark.parametrize(
+        "meta_data",
+        [
+            # Direct usage field
+            {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}},
+            # Nested metadata
+            {"metadata": {"usage": {"total_tokens": 100}}},
+            # Flat token fields
+            {"total_tokens": 50, "prompt_tokens": 30, "completion_tokens": 20},
+            # With price info
+            {
+                "usage": {
+                    "total_tokens": 150,
+                    "total_price": "0.002",
+                    "currency": "EUR",
+                }
+            },
+            # Deep nested
+            {"level1": {"level2": {"usage": {"total_tokens": 200}}}},
+        ],
+    )
+    def test_various_meta_formats(self, meta_data) -> None:
+        """Test that various meta formats are correctly parsed."""
+        result = CallToolResult(content=[], _meta=meta_data)
+        usage = MCPTool._derive_usage_from_result(result)
+
+        assert isinstance(usage, LLMUsage)
+        # Should have at least some usage data
+        if meta_data.get("usage", {}).get("total_tokens") or meta_data.get("total_tokens"):
+            expected_total = (
+                meta_data.get("usage", {}).get("total_tokens")
+                or meta_data.get("total_tokens")
+                or meta_data.get("metadata", {}).get("usage", {}).get("total_tokens")
+                or meta_data.get("level1", {}).get("level2", {}).get("usage", {}).get("total_tokens")
+            )
+            if expected_total:
+                assert usage.total_tokens == expected_total