Remove stale catalog cache from CodeMode execute

This commit is contained in:
Jeremiah Lowin 2026-03-03 16:29:18 -05:00
commit f923b671ce
2 changed files with 28 additions and 9 deletions

View file

@ -541,16 +541,9 @@ class CodeMode(CatalogTransform):
ctx: Context = None, # type: ignore[assignment]
) -> Any:
"""Execute tool calls using Python code."""
cached_tools: Sequence[Tool] | None = None
async def _get_cached_tools() -> Sequence[Tool]:
nonlocal cached_tools
if cached_tools is None:
cached_tools = await transform.get_tool_catalog(ctx)
return cached_tools
async def call_tool(tool_name: str, params: dict[str, Any]) -> Any:
backend_tools = await _get_cached_tools()
backend_tools = await transform.get_tool_catalog(ctx)
tool = transform._find_tool(tool_name, backend_tools)
if tool is None:
raise NotFoundError(f"Unknown tool: {tool_name}")

View file

@ -5,7 +5,7 @@ from typing import Any
import pytest
from mcp.types import ImageContent, TextContent
from fastmcp import FastMCP
from fastmcp import Client, FastMCP
from fastmcp.exceptions import ToolError
from fastmcp.experimental.transforms.code_mode import (
CodeMode,
@ -507,6 +507,32 @@ async def test_code_mode_search_respects_disabled_tool_visibility() -> None:
assert "secret" not in text or "No tools" in text
async def test_code_mode_execute_sees_mid_run_visibility_changes() -> None:
"""Unlocking a tool mid-execution makes it callable in the same run."""
mcp = FastMCP("CodeMode Unlock")
@mcp.tool
async def unlock(ctx: Context) -> str:
await ctx.enable_components(names={"secret"}, components={"tool"})
return "unlocked"
@mcp.tool
async def secret() -> str:
return "secret-ok"
mcp.disable(names={"secret"}, components={"tool"})
mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider()))
async with Client(mcp) as client:
result = await client.call_tool(
"execute",
{
"code": "await call_tool('unlock', {})\nreturn await call_tool('secret', {})"
},
)
assert result.data == {"result": "secret-ok"}
async def test_code_mode_execute_respects_tool_auth() -> None:
mcp = FastMCP("CodeMode Auth")