Recommended way to use InjectedState with MCP tools loaded via streamable-http server

[jonas-debeuk]

Hello!

I'm using langgraph to build an agent, which has access to tools defined in our backend server, via an MCP server serving streamable-http.

I'm trying to figure out a way to programmatically set a tenant_id, to ensure no risk of data leaking between users if the agent makes a mistake or is asked to fetch data for a diff tenant.

What is the recommended way to achieve this? I'm at a bit of a loss as I can't see how to specify the state injection into the tool definition.

here's what I have so far, but in _create_async_tool_func I'm unsure how or where to specify that tenant_id, if passed, should be an injectedstate arg.

The tool_info.inputSchema i pass for the schema is a dict, and from the langgraph code it looks like if it's in this format theres no way to specify tenant_id as an injectedstate arg?

Any help very appreciated!

I've taken a look at this thread (langchain-ai/langchain#24906) but it seems the result was inconclusive.

import asyncio
import copy
from typing import (
    Any,
    Dict,
    List,
)

from langchain_core.tools import BaseTool, StructuredTool
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client

from .logger import get_logger
from .settings import settings

logger = get_logger(__name__)


def _create_async_tool_func(tool_name: str, headers: Dict[str, str]):
    """Create an async tool function with tenant_id validation.

    Args:
        tool_name: Name of the MCP tool
        headers: HTTP headers for authentication

    Returns:
        Async function that can call the MCP tool with validation
    """

    async def async_tool_func(**kwargs):
        async with streamablehttp_client(url=settings.prediko_inventory_mcp_url, headers=headers) as tool_streams:
            tool_read_stream, tool_write_stream, _ = tool_streams
            async with ClientSession(tool_read_stream, tool_write_stream) as tool_session:
                await tool_session.initialize()
                response = await tool_session.call_tool(tool_name, kwargs)
                return response.content

    return async_tool_func


def _create_mcp_tool(tool_info, headers: Dict[str, str]) -> StructuredTool:
    """Create a LangChain StructuredTool from MCP tool info.

    Args:
        tool_info: MCP tool information object
        headers: HTTP headers for authentication

    Returns:
        LangChain StructuredTool instance
    """
    return StructuredTool.from_function(
        coroutine=_create_async_tool_func(tool_info.name, headers),
        name=tool_info.name,
        description=tool_info.description or f"Tool: {tool_info.name}",
        args_schema=tool_info.inputSchema,
    )


async def setup_prediko_tools() -> List[BaseTool]:
    """Setup Prediko inventory tools via native MCP SDK."""
    # Prepare headers with access token
    headers = {}
    if hasattr(settings, "prediko_api_token") and settings.prediko_api_token:
        headers["access_token"] = settings.prediko_api_token

    try:
        tools = []

        # Connect using streamable HTTP client
        async with streamablehttp_client(url=settings.prediko_inventory_mcp_url, headers=headers) as streams:
            read_stream, write_stream, _ = streams
            async with ClientSession(read_stream, write_stream) as session:
                # Initialize the session
                await session.initialize()

                # List available tools from MCP server
                result = await session.list_tools()

                for tool_info in result.tools:
                    tool = _create_mcp_tool(tool_info, headers)
                    tools.append(tool)

        return tools

    except Exception as e:
        logger.warning(f"Failed to connect to Prediko MCP server: {e}")
        return []

[jonas-debeuk]

For now I think I have a workaround (thanks to claude code)

1. remove the tenant_id from the args schema properties dict if it exists (and required dict also if in there)
2. when calling the tool, add a wrapper which re-injects this property according to what I have in the Agent state
   Feels a little hacky but I think this is achieving the intended result