Don't use tokens to create a Service in local tools (#64)

Instead serialize all of the auth fields and propagate them
to the local tools.

I also noted during testing that on Cloud, that token auth needs
to be enabled explicitly and without doing so the authorization/tokens
fails. With this approach we use identical credentials as the Service,
thus eliminate this problem.

Author: mateusz834

splunklib/ai/registry.py

@@ -32,8 +32,8 @@
 from mcp.server.lowlevel import Server
 from pydantic import TypeAdapter
 
-from splunklib.binding import _spliturl
-from splunklib.client import Service, connect
+from splunklib.ai.serialized_service import SerializedService
+from splunklib.client import Service
 
 
 def _normalize_logger_level(levelno: int) -> int:
@@ -147,8 +147,7 @@ class _ToolContextParams:
     in this class name i.e. internal class).
     """
 
-    management_url: str | None
-    management_token: str | None
+    service: SerializedService | None
     logger: Logger
 
 
@@ -176,21 +175,13 @@ def service(self) -> Service:
         if self._service is not None:
             return self._service
 
-        assert all((self._params.management_url, self._params.management_token)), (
-            "Invalid tool invocation, missing management_url and/or management_token"
+        assert self._params.service is not None, (
+            "Invalid tool invocation, missing serialized service details"
         )
 
-        scheme, host, port, path = _spliturl(self._params.management_url)
-        s = connect(
-            scheme=scheme,
-            host=host,
-            port=port,
-            path=path,
-            token=self._params.management_token,
-            autologin=True,
-        )
-        self._service = s
-        return s
+        # TODO: Shouldn't this function be async and this use asyncio.to_thread()?
+        self._service = self._params.service.connect()
+        return self._service
 
     @property
     def logger(self) -> Logger:
@@ -281,20 +272,19 @@ async def _call_tool(
                 logger.setLevel(_min_logging_level(self._logging_level))
                 logger.addHandler(handler)
 
-                management_url: str | None = None
-                management_token: str | None = None
+                service: SerializedService | None = None
 
                 meta = req_ctx.meta
                 if meta is not None:
                     splunk_meta = meta.model_dump().get("splunk")
                     if splunk_meta is not None:
-                        management_url = splunk_meta.get("management_url")
-                        management_token = splunk_meta.get("management_token")
+                        service = SerializedService.model_validate(
+                            splunk_meta.get("service")
+                        )
 
                 ctx = ToolContext(
                     params=_ToolContextParams(
-                        management_url=management_url,
-                        management_token=management_token,
+                        service=service,
                         logger=logger,
                     )
                 )

splunklib/ai/serialized_service.py

@@ -0,0 +1,62 @@
+# Copyright © 2011-2026 Splunk, Inc.
+#
+# 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 typing import Self
+
+from pydantic import BaseModel
+
+from splunklib.binding import _spliturl
+from splunklib.client import Service, connect
+
+
+class SerializedService(BaseModel):
+    management_url: str = ""
+    username: str | None = None
+    password: str | None = None
+    token: str | None = None
+    bearer_token: str | None = None
+    auth_cookies: dict[str, str] | None = None
+
+    @classmethod
+    def from_service(cls, service: Service) -> Self:
+        return cls(
+            management_url=f"{service.scheme}://{service.host}:{service.port}",  # pyright: ignore[reportUnknownMemberType]
+            username=service.username if service.username else None,  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
+            password=service.password if service.password else None,  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
+            token=service.token if isinstance(service.token, str) else None,  # pyright: ignore[reportUnknownMemberType, reportArgumentType]
+            bearer_token=service.bearerToken if service.bearerToken else None,  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
+            auth_cookies=(
+                service.get_cookies() if len(service.get_cookies()) != 0 else None  # pyright: ignore[reportUnknownArgumentType]
+            ),
+        )
+
+    def connect(self) -> Service:
+        scheme, host, port, path = _spliturl(self.management_url)  # pyright: ignore[reportUnknownVariableType]
+        return connect(
+            scheme=scheme,  # pyright: ignore[reportUnknownArgumentType]
+            host=host,  # pyright: ignore[reportUnknownArgumentType]
+            port=port,
+            path=path,
+            username=self.username if self.username else None,
+            password=self.password if self.password else None,
+            token=self.token if self.token else None,
+            splunkToken=self.bearer_token if self.bearer_token else None,
+            cookie="; ".join(
+                f"{key}={self.auth_cookies[key]}" for key in self.auth_cookies
+            )
+            if self.auth_cookies
+            else None,
+            autologin=True,
+        )

splunklib/ai/tools.py

@@ -23,6 +23,7 @@
 from mcp.types import Tool as MCPTool
 from pydantic import BaseModel
 
+from splunklib.ai.serialized_service import SerializedService
 from splunklib.binding import HTTPError
 from splunklib.client import Service
 from splunklib.ai.registry import _map_logger_to_mcp_logging_level
@@ -137,8 +138,7 @@ async def __call__(
 @dataclass
 class LocalCfg:
     tools_path: str
-    management_url: str
-    token: str
+    service: SerializedService
 
 
 @dataclass
@@ -256,8 +256,7 @@ async def call_tool(
                         # Provide access to the splunk instance in local tools.
                         # No need to do anything special for remote tools, since
                         # these tools are already authenticated with the token.
-                        "management_url": cfg.management_url,
-                        "management_token": cfg.token,
+                        "service": cfg.service.model_dump(),
                         # Currently we don't need to send the trace_id and app_id to local tools, since
                         # that is only really needed to correlate logs, but for local tools we know
                         # that logs coming from the local tool registry are already reladed to this
@@ -342,30 +341,6 @@ class ResponseBody(BaseModel):
     return body.entry[0].content.username
 
 
-def _get_splunk_token_for_mcp(service: Service) -> str:
-    res = service.post(
-        path_segment="authorization/tokens",
-        name=_get_splunk_username(service),
-        audience="mcp",
-        type="ephemeral",
-        output_mode="json",
-    )
-
-    class Content(BaseModel):
-        token: str
-
-    class Entry(BaseModel):
-        content: Content
-
-    class ResponseBody(BaseModel):
-        entry: list[Entry]
-
-    body = ResponseBody.model_validate_json(str(res.body))
-    if len(body.entry) == 0:
-        return ""
-    return body.entry[0].content.token
-
-
 def _get_mcp_token(service: Service) -> str | None:
     try:
         res = service.get(
@@ -401,7 +376,6 @@ async def load_mcp_tools(
 
     management_url = f"{service.scheme}://{service.host}:{service.port}"
     mcp_url = f"{management_url}/services/mcp"
-    token = await asyncio.to_thread(lambda: _get_splunk_token_for_mcp(service))
 
     mcp_token = await asyncio.to_thread(lambda: _get_mcp_token(service))
     if mcp_token is not None:
@@ -425,11 +399,7 @@ async def load_mcp_tools(
         local_tools = await _load_tools(
             LocalCfg(
                 tools_path=local_tools_path,
-                management_url=management_url,
-                # TODO: Is this right? I think we should do this differentlly and either serialize
-                # the Service auth fields and send them or generate a separate token, that does not have
-                # the "mcp" audience set.
-                token=token,
+                service=SerializedService.from_service(service),
             ),
             logger,
         )

tests/integration/ai/test_agent_mcp_tools.py

@@ -2,6 +2,7 @@
 import contextlib
 from dataclasses import asdict, dataclass
 import logging
+import json
 import os
 import socket
 from typing import Annotated
@@ -23,7 +24,6 @@
 from splunklib.ai.messages import HumanMessage, ToolMessage
 from splunklib.ai.tool_filtering import ToolFilters
 from splunklib.ai.tools import (
-    _get_splunk_token_for_mcp,
     _get_splunk_username,
     locate_app,
 )
@@ -140,21 +140,30 @@ async def test_agent_filtering_tools(self) -> None:
             assert tool_names == ["test_tool_1", "test_tool_2", "test_tool_4"]
 
 
-class TestSplunkToken(testlib.SDKTestCase):
+class TestSplunkGetUsername(testlib.SDKTestCase):
+    def get_splunk_bearer_token(self) -> str:
+        res = self.service.post(
+            path_segment="authorization/tokens",
+            name=self.service.username,
+            audience="test",
+            type="ephemeral",
+            output_mode="json",
+        )
+        token = json.loads(str(res.body))["entry"][0]["content"]["token"]
+        return token
+
     def test_get_splunk_username(self) -> None:
         self.assertTrue(
-            self.service.username is not None and self.service.username != ""
+            self.service.username and self.service.password
         )  # our CI logs-in with username and password.
 
         self.assertEqual(_get_splunk_username(self.service), self.service.username)
 
-        token = _get_splunk_token_for_mcp(self.service)
-
         service = connect(
             scheme=self.service.scheme,
             host=self.service.host,
             port=self.service.port,
-            token=token,
+            token=self.get_splunk_bearer_token(),
         )
 
         self.assertEqual(_get_splunk_username(service), self.service.username)
@@ -176,28 +185,6 @@ def test_locate_app(self) -> None:
 AUTH_TOKEN = "foobarbaz"
 
 
-async def tokens_handler(request: Request) -> Response:
-    class Content(BaseModel):
-        token: str
-
-    class Entry(BaseModel):
-        content: Content
-
-    class ResponseBody(BaseModel):
-        entry: list[Entry]
-
-    body = ResponseBody(
-        entry=[
-            Entry(content=Content(token=AUTH_TOKEN)),
-        ]
-    )
-
-    return JSONResponse(
-        content=body.model_dump(),
-        status_code=200,
-    )
-
-
 async def mcp_token_handler(_: Request) -> Response:
     return JSONResponse(
         content={"token": AUTH_TOKEN},
@@ -276,11 +263,6 @@ async def dispatch(self, request: Request, call_next):
                         mcp_token_handler,
                         methods=["GET"],
                     ),
-                    Route(
-                        "/services/authorization/tokens",
-                        tokens_handler,
-                        methods=["POST"],
-                    ),
                 ],
                 lifespan=lifespan,
                 middleware=[Middleware(MCPMiddleware)],
@@ -344,13 +326,7 @@ async def test_remote_tools_mcp_app_unavail(self):
 
         async with run_http_server(
             Starlette(
-                routes=[
-                    Route(
-                        "/services/authorization/tokens",
-                        tokens_handler,
-                        methods=["POST"],
-                    ),
-                ],
+                routes=[],
             )
         ) as (host, port):
             service = await asyncio.to_thread(
@@ -420,11 +396,6 @@ async def lifespan(app: Starlette):
                         mcp_token_handler,
                         methods=["GET"],
                     ),
-                    Route(
-                        "/services/authorization/tokens",
-                        tokens_handler,
-                        methods=["POST"],
-                    ),
                 ],
                 lifespan=lifespan,
             )
@@ -521,11 +492,6 @@ async def lifespan(app: Starlette):
                         mcp_token_handler,
                         methods=["GET"],
                     ),
-                    Route(
-                        "/services/authorization/tokens",
-                        tokens_handler,
-                        methods=["POST"],
-                    ),
                 ],
                 lifespan=lifespan,
             )