Use services/mcp_token for retrieving auth token for MCP (#63)

Splunk MCP Server App stopped supporting normal splunk tokens.
And requires calling this endpoint to retrieve the token for
use in auth to the services/mcp endpoint.

Additionally during testing I have noticed few timeouts during tool
calls. The default MCP timeouts are higher, but since we override the
default httpx client we do not get these, so lets increase these.

Author: mateusz834

splunklib/ai/tools.py

@@ -23,6 +23,7 @@
 from mcp.types import Tool as MCPTool
 from pydantic import BaseModel
 
+from splunklib.binding import HTTPError
 from splunklib.client import Service
 from splunklib.ai.registry import _map_logger_to_mcp_logging_level
 
@@ -175,6 +176,11 @@ async def _connect_local_mcp(cfg: LocalCfg, logger: _MCPLoggingHandler | None =
             yield session
 
 
+# Based on streamable_http_client defaults, when http_client is usnet.
+_MCP_DEFAULT_TIMEOUT = 30.0  # General operations (seconds)
+_MCP_DEFAULT_SSE_READ_TIMEOUT = 300.0  # SSE streams - 5 minutes (seconds)
+
+
 @asynccontextmanager
 async def _connect_remote_mcp(cfg: RemoteCfg):
     async with streamable_http_client(
@@ -187,6 +193,9 @@ async def _connect_remote_mcp(cfg: RemoteCfg):
             auth=_MCPAuth(f"Bearer {cfg.token}"),
             verify=False,
             follow_redirects=True,
+            timeout=httpx.Timeout(
+                _MCP_DEFAULT_TIMEOUT, read=_MCP_DEFAULT_SSE_READ_TIMEOUT
+            ),
         ),
     ) as (read, write, _):
         async with ClientSession(read, write) as session:
@@ -357,6 +366,24 @@ class ResponseBody(BaseModel):
     return body.entry[0].content.token
 
 
+def _get_mcp_token(service: Service) -> str | None:
+    try:
+        res = service.get(
+            path_segment="mcp_token",
+            username=_get_splunk_username(service),
+            output_mode="json",
+        )
+    except HTTPError as e:
+        if e.status == 404:
+            return None
+        raise
+
+    class ResponseBody(BaseModel):
+        token: str
+
+    return ResponseBody.model_validate_json(str(res.body)).token
+
+
 async def _load_tools(cfg: LocalCfg | RemoteCfg, logger: logging.Logger) -> list[Tool]:
     tools = await _list_all_tools(cfg)
     return [_convert_mcp_tool(logger, cfg, tool) for tool in tools]
@@ -376,13 +403,16 @@ async def load_mcp_tools(
     mcp_url = f"{management_url}/services/mcp"
     token = await asyncio.to_thread(lambda: _get_splunk_token_for_mcp(service))
 
-    # Load remote MCP tools, only if the MCP server App is available.
-    client = httpx.AsyncClient(auth=_MCPAuth(f"Bearer {token}"), verify=False)
-    res = await client.get(mcp_url)
-    if res.status_code != 404:
+    mcp_token = await asyncio.to_thread(lambda: _get_mcp_token(service))
+    if mcp_token is not None:
         logger.debug("Splunk MCP Server App detected - loading remote tools")
         remote_tools = await _load_tools(
-            RemoteCfg(mcp_url=mcp_url, token=token, app_id=app_id, trace_id=trace_id),
+            RemoteCfg(
+                mcp_url=mcp_url,
+                token=mcp_token,
+                app_id=app_id,
+                trace_id=trace_id,
+            ),
             logger,
         )
         logger.debug(

tests/integration/ai/test_agent_mcp_tools.py

@@ -1,6 +1,7 @@
 import asyncio
 import contextlib
 from dataclasses import asdict, dataclass
+import logging
 import os
 import socket
 from typing import Annotated
@@ -197,6 +198,13 @@ class ResponseBody(BaseModel):
     )
 
 
+async def mcp_token_handler(_: Request) -> Response:
+    return JSONResponse(
+        content={"token": AUTH_TOKEN},
+        status_code=200,
+    )
+
+
 class TestRemoteTools(AITestCase):
     @patch(
         "splunklib.ai.agent._testing_local_tools_path",
@@ -263,6 +271,11 @@ async def dispatch(self, request: Request, call_next):
             Starlette(
                 routes=[
                     Mount("/services/mcp", app=mcp.streamable_http_app()),
+                    Route(
+                        "/services/mcp_token",
+                        mcp_token_handler,
+                        methods=["GET"],
+                    ),
                     Route(
                         "/services/authorization/tokens",
                         tokens_handler,
@@ -402,6 +415,11 @@ async def lifespan(app: Starlette):
             Starlette(
                 routes=[
                     Mount("/services/mcp", app=mcp.streamable_http_app()),
+                    Route(
+                        "/services/mcp_token",
+                        mcp_token_handler,
+                        methods=["GET"],
+                    ),
                     Route(
                         "/services/authorization/tokens",
                         tokens_handler,
@@ -498,6 +516,11 @@ async def lifespan(app: Starlette):
             Starlette(
                 routes=[
                     Mount("/services/mcp", app=mcp.streamable_http_app()),
+                    Route(
+                        "/services/mcp_token",
+                        mcp_token_handler,
+                        methods=["GET"],
+                    ),
                     Route(
                         "/services/authorization/tokens",
                         tokens_handler,
@@ -551,6 +574,50 @@ async def lifespan(app: Starlette):
                 response = result.messages[-1].content
                 assert "31.5" in response, "Invalid LLM response"
 
+    @patch(
+        "splunklib.ai.agent._testing_local_tools_path",
+        os.path.join(
+            os.path.dirname(__file__),
+            "testdata",
+            "non_existent.py",
+        ),
+    )
+    @patch("splunklib.ai.agent._testing_app_id", "app_id")
+    @pytest.mark.asyncio
+    async def test_splunk_mcp_server_app(self) -> None:
+        # Skip if the langchain_openai package is not installed
+        pytest.importorskip("langchain_openai")
+
+        # TODO: Remove this test once we have an E2E with Splunk MCP Server app.
+
+        self.skipTest("manual test")
+
+        logger = logging.getLogger("test")
+        logger.setLevel(logging.DEBUG)
+
+        service = connect(
+            port=8090,
+            host="localhost",
+            username="admin",
+            password="",
+            autologin=True,
+        )
+
+        async with Agent(
+            model=(await self.model()),
+            system_prompt="You must use the available tools to perform requested operations",
+            service=service,
+            use_mcp_tools=True,
+            logger=logger,
+        ) as agent:
+            for tool in agent.tools:
+                if tool.name == "splunk_get_indexes":
+                    result = await tool.func()
+                    assert len(result.structured_content["results"]) != 0
+                    return
+
+            assert False, "tool splunk_get_indexes not found"
+
 
 @contextlib.asynccontextmanager
 async def run_http_server(app: Starlette):