DX-118683: Access tools from Dremio repository (#99)

* First commit

* Fixed tests

* Fixing tests

* Fixing tests

* Fixed minor errors

* Fix small problems

* Wired tools to MCP

* Finished review

* Fixed feedback

* Addressed feedback

* Fix test failure

* Fix test

* Addressed feedback

* Cleanup

* Addressed feedback

* Added flag

* Addressed feedback

Author: AfonsoP1nto

src/dremioai/api/dremio/ai_tools.py

@@ -0,0 +1,93 @@
+#
+#  Copyright (C) 2017-2025 Dremio Corporation
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+from pydantic import BaseModel, Field, ConfigDict
+from typing import Dict, List, Optional, Any
+from urllib.parse import quote
+
+from aiohttp import ClientResponseError
+from dremioai.api.transport import DremioAsyncHttpClient as AsyncHttpClient
+from dremioai.log import logger
+
+log = logger(__name__)
+
+
+class AiTool(BaseModel):
+
+    name: str
+    description: Optional[str] = None
+    input_schema: Dict[str, Any] = Field(
+        default_factory=lambda: {"type": "object"}, alias="inputSchema"
+    )
+    model_config = ConfigDict(extra="allow", populate_by_name=True)
+
+
+class ListToolsResponse(BaseModel):
+    tools: List[AiTool] = Field(default_factory=list)
+    error: Optional[str] = None
+
+    def __bool__(self):
+        return self.error is None
+
+
+class InvokeToolResponse(BaseModel):
+
+    result: Optional[Any] = None
+    error: Optional[str] = None
+
+    def __bool__(self):
+        return self.error is None
+
+    @property
+    def is_empty(self) -> bool:
+        """True when the response carries neither a result nor an error.
+
+        This can happen when Dremio returns a 200 with an empty body for a
+        void tool.  Callers may choose to treat this as a successful no-op.
+        """
+        return self.result is None and self.error is None
+
+
+async def list_tools() -> ListToolsResponse:
+    try:
+        client = AsyncHttpClient()
+        return await client.get(
+            "/api/v4/ai/tools",
+            deser=ListToolsResponse,
+        )
+    except ClientResponseError as e:
+        log.exception("Failed to list AI tools")
+        return ListToolsResponse(error=f"HTTP {e.status} {e.message}")
+    except Exception:
+        log.exception("Failed to list AI tools")
+        return ListToolsResponse(error="Unexpected error listing AI tools")
+
+
+async def invoke_tool(tool_name: str, args: Dict[str, Any]) -> InvokeToolResponse:
+    safe_name = quote(tool_name, safe="")
+    try:
+        client = AsyncHttpClient()
+        return await client.post(
+            f"/api/v4/ai/tools/{safe_name}:invoke",
+            body={"args": args},
+            deser=InvokeToolResponse,
+        )
+    except ClientResponseError as e:
+        log.exception("Failed to invoke AI tool '%s'", tool_name)
+        return InvokeToolResponse(error=f"HTTP {e.status} {e.message}")
+    except Exception:
+        log.exception("Failed to invoke AI tool '%s'", tool_name)
+        return InvokeToolResponse(error=f"Unexpected error invoking tool '{tool_name}'")

src/dremioai/config/settings.py

@@ -274,6 +274,10 @@ class Dremio(FlagAwareModel):
     wlm: Optional[Wlm] = None
     api: Optional[ApiSettings] = Field(default_factory=ApiSettings)
     metrics: Optional[Metrics] = None
+    enable_remote_tools: Optional[bool] = Field(
+        default=False,
+        description="Enable dynamic registration of remote tools from Dremio's Java-side tool registry",
+    )
 
     @field_serializer("raw_pat")
     def serialize_pat(self, pat: str):

src/dremioai/servers/mcp.py

@@ -34,7 +34,7 @@
 from click import Choice
 from mcp.cli.claude import get_claude_config_path
 from mcp.server.auth.json_response import PydanticJSONResponse
-from mcp.server.auth.middleware.auth_context import AuthContextMiddleware
+from mcp.server.auth.middleware.auth_context import AuthContextMiddleware, get_access_token
 from mcp.server.auth.middleware.bearer_auth import BearerAuthBackend
 from mcp.server.auth.provider import AccessToken, TokenVerifier
 from mcp.server.fastmcp import FastMCP
@@ -59,6 +59,7 @@
 from dremioai.config import settings
 from dremioai.config.feature_flags import FeatureFlagManager
 from dremioai.metrics.registry import get_metrics_app
+from dremioai.metrics.tool_metrics import invocation_counter, invocation_duration
 from dremioai.servers.jwks_verifier import JWKSVerifier, TokenExpiredError
 from dremioai.tools import tools
 from dremioai.tools.tools import ProjectIdMiddleware

src/dremioai/tools/tools.py

@@ -44,14 +44,15 @@
 
 import pandas as pd
 import numpy as np
-from dremioai.api.dremio import sql, usage, search
+from dremioai.api.dremio import sql, usage, search, ai_tools
 from dremioai.config import settings
 from dremioai.config.tools import ToolType
 from dremioai.api.prometheus import vm
 from dremioai.api.dremio.catalog import get_schema, get_lineage, get_descriptions
 from dremioai.api.util import run_in_parallel
 from csv import reader
 from io import StringIO
+import json
 from sqlglot import parse_one
 from sqlglot import expressions
 from mcp.server.auth.middleware.auth_context import get_access_token
@@ -537,6 +538,47 @@ async def invoke(self, query: str) -> Dict[str, Any]:
         return {"results": res.to_dict(orient="records")}
 
 
+class DiscoverDynamicTools(Tools):
+    For: ClassVar[Annotated[ToolType, ToolType.FOR_SELF | ToolType.FOR_DATA_PATTERNS]]
+
+    @secured
+    @with_metrics
+    async def invoke(self) -> str:
+        """Discover additional tools available from the Dremio server.
+        Call this tool to get a list of dynamically available tools with their
+        names, descriptions, and input schemas."""
+        if not settings.instance().dremio.get("enable_remote_tools"):
+            return "Remote tools are not enabled."
+        result = await ai_tools.list_tools()
+        return result.model_dump_json()
+
+
+class CallDynamicTool(Tools):
+    For: ClassVar[Annotated[ToolType, ToolType.FOR_SELF | ToolType.FOR_DATA_PATTERNS]]
+
+    @secured
+    @with_metrics
+    async def invoke(self, tool_name: str, tool_arguments: Union[str, dict]) -> str:
+        """Invoke a dynamically discovered tool on the Dremio server.
+
+        Args:
+            tool_name: The name of the tool to invoke, as returned by DiscoverDynamicTools.
+            tool_arguments: The arguments to pass to the tool, either as a JSON string or a dict.
+        """
+        if not settings.instance().dremio.get("enable_remote_tools"):
+            return "Remote tools are not enabled."
+        if isinstance(tool_arguments, str):
+            try:
+                args = json.loads(tool_arguments)
+            except json.JSONDecodeError as exc:
+                return f"Invalid JSON in tool_arguments: {exc}"
+        else:
+            args = tool_arguments
+
+        result = await ai_tools.invoke_tool(tool_name, args)
+        return result.model_dump_json(exclude_none=True)
+
+
 def _subclasses(cls):
     for sub in cls.__subclasses__():
         yield from _subclasses(sub)

tests/api/dremio/test_ai_tools.py

@@ -0,0 +1,174 @@
+#
+#  Copyright (C) 2017-2025 Dremio Corporation
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+import pytest
+from dremioai.api.dremio.ai_tools import (
+    AiTool,
+    InvokeToolResponse,
+    list_tools,
+    invoke_tool,
+)
+from mocks.http_mock import HttpMockFramework
+
+
+# --- list_tools tests ---
+
+@pytest.mark.asyncio
+async def test_list_tools_returns_tools(mock_settings_instance):
+    with HttpMockFramework() as mock:
+        mock.load_mock_data(r"/api/v4/ai/tools$", "ai_tools/list_tools.json")
+        result = await list_tools()
+    assert bool(result)
+    assert len(result.tools) == 3
+    names = [t.name for t in result.tools]
+    assert "runSql" in names
+    assert "getTableOrViewSchema" in names
+    assert "listEngines" in names
+
+
+@pytest.mark.asyncio
+async def test_list_tools_returns_input_schema(mock_settings_instance):
+    with HttpMockFramework() as mock:
+        mock.load_mock_data(r"/api/v4/ai/tools$", "ai_tools/list_tools.json")
+        result = await list_tools()
+    run_sql = next(t for t in result.tools if t.name == "runSql")
+    assert run_sql.input_schema["type"] == "object"
+    assert "sqlText" in run_sql.input_schema["properties"]
+
+
+@pytest.mark.asyncio
+async def test_list_tools_empty_registry(mock_settings_instance):
+    with HttpMockFramework() as mock:
+        mock.add_mock_response(r"/api/v4/ai/tools$", {"tools": []})
+        result = await list_tools()
+    assert result.tools == []
+    assert bool(result)
+
+
+# --- invoke_tool tests ---
+
+@pytest.mark.asyncio
+async def test_invoke_tool_success(mock_settings_instance):
+    with HttpMockFramework() as mock:
+        mock.load_mock_data(r"/api/v4/ai/tools/runSql:invoke$", "ai_tools/invoke_result.json")
+        result = await invoke_tool("runSql", {"sqlText": "SELECT 1"})
+    assert bool(result)
+    assert result.result["columns"] == ["id", "name"]
+    assert result.error is None
+
+
+@pytest.mark.asyncio
+async def test_invoke_tool_error_response(mock_settings_instance):
+    with HttpMockFramework() as mock:
+        mock.load_mock_data(r"/api/v4/ai/tools/unknownTool:invoke$", "ai_tools/invoke_error.json")
+        result = await invoke_tool("unknownTool", {})
+    assert result.error is not None
+    assert "not found" in result.error
+    assert result.result is None
+
+
+# --- Pydantic model unit tests (no HTTP) ---
+
+def test_ai_tool_model_validation():
+    raw = {
+        "name": "runSql",
+        "description": "runSql",
+        "inputSchema": {
+            "type": "object",
+            "properties": {"sqlText": {"type": "string"}},
+            "required": ["sqlText"],
+        },
+    }
+    tool = AiTool.model_validate(raw)
+    assert tool.name == "runSql"
+    assert tool.input_schema["type"] == "object"
+    assert tool.input_schema["required"] == ["sqlText"]
+
+
+def test_ai_tool_model_minimal_schema():
+    """Tools with an empty inputSchema (e.g. listEngines) should deserialize cleanly."""
+    raw = {"name": "listEngines", "description": "listEngines", "inputSchema": {"type": "object"}}
+    tool = AiTool.model_validate(raw)
+    assert tool.name == "listEngines"
+    assert tool.input_schema == {"type": "object"}
+
+
+def test_invoke_tool_response_succeeded():
+    resp = InvokeToolResponse.model_validate({"result": {"sql": "SELECT 1"}})
+    assert bool(resp) is True
+    assert resp.result == {"sql": "SELECT 1"}
+    assert resp.error is None
+
+
+def test_invoke_tool_response_failed():
+    resp = InvokeToolResponse.model_validate({"error": "Tool not found"})
+    assert bool(resp) is False
+    assert resp.result is None
+    assert resp.error == "Tool not found"
+
+
+def test_invoke_tool_response_empty():
+    resp = InvokeToolResponse.model_validate({})
+    assert bool(resp) is True
+    assert resp.result is None
+    assert resp.error is None
+    assert resp.is_empty is True
+
+
+def test_invoke_tool_response_is_empty_false_when_result():
+    """is_empty should be False when a result is present."""
+    resp = InvokeToolResponse.model_validate({"result": {"sql": "SELECT 1"}})
+    assert resp.is_empty is False
+
+
+def test_invoke_tool_response_is_empty_false_when_error():
+    """is_empty should be False when an error is present."""
+    resp = InvokeToolResponse.model_validate({"error": "Tool not found"})
+    assert resp.is_empty is False
+
+
+# --- HTTP error scenario tests ---
+
+@pytest.mark.asyncio
+async def test_list_tools_http_error(mock_settings_instance):
+    """list_tools should return a response with error set on HTTP 4xx/5xx."""
+    with HttpMockFramework() as mock:
+        mock.add_mock_response(r"/api/v4/ai/tools$", {"error": "Unauthorized"}, status=401)
+        result = await list_tools()
+    assert not bool(result)
+    assert result.error is not None
+    assert "401" in result.error
+
+
+@pytest.mark.asyncio
+async def test_invoke_tool_http_error(mock_settings_instance):
+    """invoke_tool should return a response with error set on HTTP 500."""
+    with HttpMockFramework() as mock:
+        mock.add_mock_response(r"/api/v4/ai/tools/runSql:invoke$", {"error": "Internal Server Error"}, status=500)
+        result = await invoke_tool("runSql", {"sqlText": "SELECT 1"})
+    assert not bool(result)
+    assert result.error is not None
+    assert "500" in result.error
+
+
+@pytest.mark.asyncio
+async def test_invoke_tool_url_encodes_name(mock_settings_instance):
+    """tool_name with special characters should be URL-encoded."""
+    with HttpMockFramework() as mock:
+        # The encoded name "my%2Ftool" should appear in the URL
+        mock.add_mock_response(r"/api/v4/ai/tools/my%2Ftool:invoke$", {"result": "ok", "error": None})
+        result = await invoke_tool("my/tool", {})
+    assert result.result == "ok"

tests/config/golden_flag_keys.yaml

@@ -6,6 +6,7 @@ flag_keys:
 - dremio.api.http_retry.max_retries
 - dremio.api.polling_interval
 - dremio.auth_issuer_uri_override
+- dremio.enable_remote_tools
 - dremio.enable_search
 - dremio.extract_org_id_from_jwt
 - dremio.jwks_cache_lifespan