Propagate trace_id & app_id to MCP (#49)

This change makes the Agent generate a trace_id during startup
and propagates it with app_id to the MCP Server App in `_meta`
fields of MCP requests and in the HTTP headers.

Author: mateusz834

splunklib/ai/agent.py

@@ -26,15 +26,12 @@
 from splunklib.ai.messages import AgentResponse, BaseMessage, OutputT
 from splunklib.ai.model import PredefinedModel
 from splunklib.ai.tool_filtering import ToolFilters, filter_tools
-from splunklib.ai.tools import (
-    Tool,
-    load_mcp_tools,
-    locate_tools_path_by_sdk_location,
-)
+from splunklib.ai.tools import Tool, build_local_tools_path, load_mcp_tools, locate_app
 from splunklib.client import Service
 
 # For testing purposes, overrides the automatically inferred tools.py path.
 _testing_local_tools_path: str | None = None
+_testing_app_id: str | None = None
 
 
 @final
@@ -149,7 +146,9 @@ async def __aenter__(self) -> Self:
             raise AssertionError("Agent is already in `async with` context")
 
         if self._use_mcp_tools:
-            self._tools = await _load_tools_from_mcp(self._service, self._tool_filters)
+            self._tools = await _load_tools_from_mcp(
+                self._service, self._tool_filters, self.trace_id
+            )
 
         backend = get_backend()
         self._impl = await backend.create_agent(self)
@@ -169,16 +168,25 @@ async def invoke(self, messages: list[BaseMessage]) -> AgentResponse[OutputT]:
 
 
 async def _load_tools_from_mcp(
-    service: Service, filters: ToolFilters | None
+    service: Service,
+    filters: ToolFilters | None,
+    trace_id: str,
 ) -> list[Tool]:
     local_tools_path = _testing_local_tools_path
+    app_id = _testing_app_id
+
     if local_tools_path is None:
-        local_tools_path = locate_tools_path_by_sdk_location()
+        app_id, app_dir = locate_app()
+        local_tools_path = build_local_tools_path(app_dir)
+
+    assert app_id is not None, (
+        "_load_tools_from_mcp was mocked, but _testing_app_id not"
+    )
 
     if not os.path.exists(local_tools_path):
         local_tools_path = None
 
-    mcp_tools = await load_mcp_tools(service, local_tools_path)
+    mcp_tools = await load_mcp_tools(service, local_tools_path, app_id, trace_id)
     if filters:
         return filter_tools(mcp_tools, filters)
 

splunklib/ai/base_agent.py

@@ -15,6 +15,7 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
+import secrets
 from typing import Generic
 
 from pydantic import BaseModel
@@ -35,6 +36,7 @@ class BaseAgent(Generic[OutputT], ABC):
     _input_schema: type[BaseModel] | None = None
     _output_schema: type[OutputT] | None = None
     _hooks: Sequence[AgentHook] | None = None
+    _trace_id: str
 
     def __init__(
         self,
@@ -57,6 +59,7 @@ def __init__(
         self._input_schema = input_schema
         self._output_schema = output_schema
         self._hooks = tuple(hooks) if hooks else ()
+        self._trace_id = secrets.token_hex(16)  # 32 Hex characters
 
     @abstractmethod
     async def invoke(self, messages: list[BaseMessage]) -> AgentResponse[OutputT]: ...
@@ -96,3 +99,7 @@ def output_schema(self) -> type[OutputT] | None:
     @property
     def hooks(self) -> Sequence[AgentHook] | None:
         return self._hooks
+
+    @property
+    def trace_id(self) -> str:
+        return self._trace_id

splunklib/ai/registry.py

@@ -35,6 +35,7 @@ class ToolContext:
 
     _management_url: str | None = None
     _management_token: str | None = None
+
     _service: Service | None = None
 
     @property

splunklib/ai/tools.py

@@ -47,9 +47,9 @@ def _splunk_home() -> str:
     return splunk_home
 
 
-def locate_tools_path_by_sdk_location(
+def locate_app(
     splunk_home: str | None = None, sdk_location_path: str = __file__
-) -> str:
+) -> tuple[str, str]:
     """
     This function returns the path to the tools file of the app, assumes that the SDK
     is vendored into the app.
@@ -76,7 +76,11 @@ def locate_tools_path_by_sdk_location(
     assert parts[0] != "." and parts[1] != ".."
 
     app_id = parts[0]
-    return os.path.join(splunk_home, "etc", "apps", app_id, "bin", TOOLS_FILENAME)
+    return (app_id, os.path.join(splunk_home, "etc", "apps", app_id))
+
+
+def build_local_tools_path(dir: str) -> str:
+    return os.path.join(dir, "bin", TOOLS_FILENAME)
 
 
 @dataclass
@@ -90,6 +94,8 @@ class LocalCfg:
 class RemoteCfg:
     mcp_url: str
     token: str
+    app_id: str
+    trace_id: str
 
 
 @asynccontextmanager
@@ -120,6 +126,10 @@ async def _connect_remote_mcp(cfg: RemoteCfg):
     async with streamable_http_client(
         url=cfg.mcp_url,
         http_client=httpx.AsyncClient(
+            headers={
+                "x-splunk-trace-id": cfg.trace_id,
+                "x-splunk-app-id": cfg.app_id,
+            },
             auth=_MCPAuth(f"Bearer {cfg.token}"),
             verify=False,
             follow_redirects=True,
@@ -174,17 +184,29 @@ def _convert_mcp_tool(
     async def call_tool(
         **arguments: dict[str, Any],
     ) -> ToolResult:
-        # 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.
         meta: dict[str, Any] | None = None
-        if isinstance(cfg, LocalCfg):
-            meta = {
-                "splunk": {
-                    "management_url": cfg.management_url,
-                    "management_token": cfg.token,
+        match cfg:
+            case LocalCfg():
+                meta = {
+                    "splunk": {
+                        # 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,
+                        # 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
+                        # agent.
+                    }
+                }
+            case RemoteCfg():
+                meta = {
+                    "splunk": {
+                        "trace_id": cfg.trace_id,
+                        "app_id": cfg.app_id,
+                    }
                 }
-            }
 
         async with _connect(cfg) as session:
             call_tool_result = await session.call_tool(
@@ -291,7 +313,9 @@ async def _load_tools(cfg: LocalCfg | RemoteCfg) -> list[Tool]:
 
 async def load_mcp_tools(
     service: Service,
-    local_tools_path: str | None = None,
+    local_tools_path: str | None,
+    app_id: str,
+    trace_id: str,
 ) -> list[Tool]:
     # TODO: Add tool.name collision between local/remote tools
     tools: list[Tool] = []
@@ -304,7 +328,9 @@ async def load_mcp_tools(
     client = httpx.AsyncClient(auth=_MCPAuth(f"Bearer {token}"), verify=False)
     res = await client.get(mcp_url)
     if res.status_code != 404:
-        remote_tools = await _load_tools(RemoteCfg(mcp_url=mcp_url, token=token))
+        remote_tools = await _load_tools(
+            RemoteCfg(mcp_url=mcp_url, token=token, app_id=app_id, trace_id=trace_id)
+        )
         tools.extend(remote_tools)
 
     if local_tools_path is not None:

tests/integration/ai/test_agent_mcp_tools.py

@@ -5,21 +5,23 @@
 from unittest.mock import patch
 
 import pytest
+from starlette.middleware import Middleware
 import uvicorn
-from mcp.server.fastmcp import FastMCP
+from mcp.server.fastmcp import Context, FastMCP
 from pydantic import BaseModel
 from starlette.applications import Starlette
 from starlette.requests import Request
 from starlette.responses import JSONResponse, Response
 from starlette.routing import Mount, Route
+from starlette.middleware.base import BaseHTTPMiddleware
 
 from splunklib.ai import Agent
 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_tools_path_by_sdk_location,
+    locate_app,
 )
 from splunklib.client import connect
 from tests import testlib
@@ -38,6 +40,7 @@ class TestTools(AITestCase):
             "weather.py",
         ),
     )
+    @patch("splunklib.ai.agent._testing_app_id", "app_id")
     async def test_tool_execution_structured_output(self) -> None:
         # Skip if the langchain_openai package is not installed
         pytest.importorskip("langchain_openai")
@@ -77,6 +80,7 @@ async def test_tool_execution_structured_output(self) -> None:
             "tool_context.py",
         ),
     )
+    @patch("splunklib.ai.agent._testing_app_id", "app_id")
     async def test_tool_execution_service_access(self) -> None:
         # Skip if the langchain_openai package is not installed
         pytest.importorskip("langchain_openai")
@@ -114,6 +118,7 @@ async def test_tool_execution_service_access(self) -> None:
         "splunklib.ai.agent._testing_local_tools_path",
         os.path.join(os.path.dirname(__file__), "testdata", "tool_filtering.py"),
     )
+    @patch("splunklib.ai.agent._testing_app_id", "app_id")
     @pytest.mark.asyncio
     async def test_agent_filtering_tools(self) -> None:
         pytest.importorskip("langchain_openai")
@@ -151,16 +156,17 @@ def test_get_splunk_username(self) -> None:
         self.assertEqual(_get_splunk_username(service), self.service.username)
 
 
-class TestToolsPathInference:
-    def test_infer_tools_path(self) -> None:
+class TestAppLocate:
+    def test_locate_app(self) -> None:
         path = os.path.join(os.path.dirname(__file__), "testdata", "app-inference")
-        got = locate_tools_path_by_sdk_location(
+        app_id, app_dir = locate_app(
             splunk_home=path,
             sdk_location_path=os.path.join(
                 path, "etc", "apps", "appname", "bin", "lib", "somefile.py"
             ),
         )
-        assert got == os.path.join(path, "etc", "apps", "appname", "bin", "tools.py")
+        assert app_id == "appname"
+        assert app_dir == os.path.join(path, "etc", "apps", "appname")
 
 
 AUTH_TOKEN = "foobarbaz"
@@ -197,14 +203,26 @@ class TestRemoteTools(AITestCase):
             "non_existent.py",
         ),
     )
+    @patch("splunklib.ai.agent._testing_app_id", "fancyapp")
     @pytest.mark.asyncio
     async def test_remote_tools(self):
         pytest.importorskip("langchain_openai")
 
         mcp = FastMCP("MCP Server", streamable_http_path="/")
 
+        trace_id: str | None = None
+        app_id: str | None = None
+
         @mcp.tool(description="Returns the current temperature in the city")
-        def temperature(city: str) -> str:
+        def temperature(ctx: Context, city: str) -> str:
+            nonlocal trace_id, app_id
+            assert trace_id is None and app_id is None
+            assert ctx.request_context.meta is not None
+            meta = ctx.request_context.meta.model_dump()
+            splunk = meta.get("splunk", {})
+            trace_id = splunk.get("trace_id")
+            app_id = splunk.get("app_id")
+
             if city == "Krakow":
                 return "31.5C"
             else:
@@ -215,6 +233,29 @@ async def lifespan(app: Starlette):
             async with mcp.session_manager.run():
                 yield
 
+        http_trace_id: str | None = None
+        http_app_id: str | None = None
+        middleware_called = False
+
+        class MCPMiddleware(BaseHTTPMiddleware):
+            async def dispatch(self, request: Request, call_next):
+                if request.url.path.startswith("/services/mcp/"):
+                    nonlocal http_trace_id, http_app_id, middleware_called
+
+                    trace_id = request.headers.get("x-splunk-trace-id")
+                    app_id = request.headers.get("x-splunk-app-id")
+
+                    # Make sure header values do not change over time.
+                    if middleware_called:
+                        assert http_trace_id == trace_id
+                        assert http_app_id == app_id
+
+                    middleware_called = True
+                    http_trace_id = trace_id
+                    http_app_id = app_id
+
+                return await call_next(request)
+
         async with run_http_server(
             Starlette(
                 routes=[
@@ -226,6 +267,7 @@ async def lifespan(app: Starlette):
                     ),
                 ],
                 lifespan=lifespan,
+                middleware=[Middleware(MCPMiddleware)],
             )
         ) as (host, port):
             service = await asyncio.to_thread(
@@ -266,6 +308,11 @@ async def lifespan(app: Starlette):
                 response = result.messages[-1].content
                 assert "31.5" in response, "Invalid LLM response"
 
+                assert trace_id == agent.trace_id
+                assert app_id == "fancyapp"
+                assert http_trace_id == agent.trace_id
+                assert http_app_id == "fancyapp"
+
     @patch(
         "splunklib.ai.agent._testing_local_tools_path",
         os.path.join(
@@ -274,6 +321,7 @@ async def lifespan(app: Starlette):
             "non_existent.py",
         ),
     )
+    @patch("splunklib.ai.agent._testing_app_id", "app_id")
     @pytest.mark.asyncio
     async def test_remote_tools_mcp_app_unavail(self):
         pytest.importorskip("langchain_openai")
@@ -326,6 +374,7 @@ async def test_remote_tools_mcp_app_unavail(self):
             "non_existent.py",
         ),
     )
+    @patch("splunklib.ai.agent._testing_app_id", "app_id")
     @pytest.mark.asyncio
     async def test_remote_tools_failure(self):
         pytest.importorskip("langchain_openai")