Clean up mode kwarg

This commit is contained in:
Jeremiah Lowin 2025-06-19 12:44:34 -04:00
commit d8235e1af7
3 changed files with 43 additions and 49 deletions

View file

@ -2,7 +2,7 @@ from __future__ import annotations as _annotations
import warnings
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any, Literal
from typing import TYPE_CHECKING, Any
from mcp import GetPromptResult
@ -46,25 +46,23 @@ class PromptManager:
"""Adds a mounted server as a source for prompts."""
self._mounted_sources.append(server)
async def _load_prompts(
self, *, mode: Literal["inventory", "protocol"]
) -> dict[str, Prompt]:
async def _load_prompts(self, *, via_server: bool = False) -> dict[str, Prompt]:
"""
The single, consolidated recursive method for fetching prompts. The 'mode'
The single, consolidated recursive method for fetching prompts. The 'via_server'
parameter determines the communication path.
- mode="inventory": Manager-to-manager path for complete, unfiltered inventory
- mode="protocol": Server-to-server path for filtered MCP requests
- via_server=False: Manager-to-manager path for complete, unfiltered inventory
- via_server=True: Server-to-server path for filtered MCP requests
"""
all_prompts: dict[str, Prompt] = {}
for mounted in self._mounted_sources:
try:
if mode == "protocol":
# PATH 2: Use the server-to-server filtered path
if via_server:
# Use the server-to-server filtered path
child_results = await mounted.server._list_prompts()
else: # mode == "inventory"
# PATH 1: Use the manager-to-manager unfiltered path
else:
# Use the manager-to-manager unfiltered path
child_results = await mounted.server._prompt_manager._list_prompts()
# The combination logic is the same for both paths
@ -104,13 +102,13 @@ class PromptManager:
"""
Gets the complete, unfiltered inventory of all prompts.
"""
return await self._load_prompts(mode="inventory")
return await self._load_prompts(via_server=False)
async def _list_prompts(self) -> list[Prompt]:
"""
Lists all prompts, applying protocol filtering.
"""
prompts_dict = await self._load_prompts(mode="protocol")
prompts_dict = await self._load_prompts(via_server=True)
return list(prompts_dict.values())
def add_prompt_from_fn(

View file

@ -5,7 +5,7 @@ from __future__ import annotations
import inspect
import warnings
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, Literal
from typing import TYPE_CHECKING, Any
from pydantic import AnyUrl
@ -63,34 +63,32 @@ class ResourceManager:
async def get_resources(self) -> dict[str, Resource]:
"""Get all registered resources, keyed by URI."""
return await self._load_resources(mode="inventory")
return await self._load_resources(via_server=False)
async def get_resource_templates(self) -> dict[str, ResourceTemplate]:
"""Get all registered templates, keyed by URI template."""
return await self._load_resource_templates(mode="inventory")
return await self._load_resource_templates(via_server=False)
async def _load_resources(
self, *, mode: Literal["inventory", "protocol"]
) -> dict[str, Resource]:
async def _load_resources(self, *, via_server: bool = False) -> dict[str, Resource]:
"""
The single, consolidated recursive method for fetching resources. The 'mode'
The single, consolidated recursive method for fetching resources. The 'via_server'
parameter determines the communication path.
- mode="inventory": Manager-to-manager path for complete, unfiltered inventory
- mode="protocol": Server-to-server path for filtered MCP requests
- via_server=False: Manager-to-manager path for complete, unfiltered inventory
- via_server=True: Server-to-server path for filtered MCP requests
"""
all_resources: dict[str, Resource] = {}
for mounted in self._mounted_sources:
try:
if mode == "protocol":
# PATH 2: Use the server-to-server filtered path
if via_server:
# Use the server-to-server filtered path
child_resources_list = await mounted.server._list_resources()
child_resources = {
resource.key: resource for resource in child_resources_list
}
else: # mode == "inventory"
# PATH 1: Use the manager-to-manager unfiltered path
else:
# Use the manager-to-manager unfiltered path
child_resources = (
await mounted.server._resource_manager.get_resources()
)
@ -120,24 +118,24 @@ class ResourceManager:
return all_resources
async def _load_resource_templates(
self, *, mode: Literal["inventory", "protocol"]
self, *, via_server: bool = False
) -> dict[str, ResourceTemplate]:
"""
The single, consolidated recursive method for fetching templates. The 'mode'
The single, consolidated recursive method for fetching templates. The 'via_server'
parameter determines the communication path.
- mode="inventory": Manager-to-manager path for complete, unfiltered inventory
- mode="protocol": Server-to-server path for filtered MCP requests
- via_server=False: Manager-to-manager path for complete, unfiltered inventory
- via_server=True: Server-to-server path for filtered MCP requests
"""
all_templates: dict[str, ResourceTemplate] = {}
for mounted in self._mounted_sources:
try:
if mode == "protocol":
# PATH 2: Use the server-to-server filtered path
if via_server:
# Use the server-to-server filtered path
child_templates = await mounted.server._list_resource_templates()
else: # mode == "inventory"
# PATH 1: Use the manager-to-manager unfiltered path
else:
# Use the manager-to-manager unfiltered path
child_templates = await mounted.server._resource_manager._list_resource_templates()
child_dict = {template.key: template for template in child_templates}
@ -169,14 +167,14 @@ class ResourceManager:
"""
Lists all resources, applying protocol filtering.
"""
resources_dict = await self._load_resources(mode="protocol")
resources_dict = await self._load_resources(via_server=True)
return list(resources_dict.values())
async def _list_resource_templates(self) -> list[ResourceTemplate]:
"""
Lists all templates, applying protocol filtering.
"""
templates_dict = await self._load_resource_templates(mode="protocol")
templates_dict = await self._load_resource_templates(via_server=True)
return list(templates_dict.values())
def add_resource_or_template_from_fn(

View file

@ -2,7 +2,7 @@ from __future__ import annotations
import warnings
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, Literal
from typing import TYPE_CHECKING, Any
from mcp.types import ToolAnnotations
@ -47,25 +47,23 @@ class ToolManager:
"""Adds a mounted server as a source for tools."""
self._mounted_sources.append(server)
async def _load_tools(
self, *, mode: Literal["inventory", "protocol"]
) -> dict[str, Tool]:
async def _load_tools(self, *, via_server: bool = False) -> dict[str, Tool]:
"""
The single, consolidated recursive method for fetching tools. The 'mode'
The single, consolidated recursive method for fetching tools. The 'via_server'
parameter determines the communication path.
- mode="inventory": Manager-to-manager path for complete, unfiltered inventory
- mode="protocol": Server-to-server path for filtered MCP requests
- via_server=False: Manager-to-manager path for complete, unfiltered inventory
- via_server=True: Server-to-server path for filtered MCP requests
"""
all_tools: dict[str, Tool] = {}
for mounted in self._mounted_sources:
try:
if mode == "protocol":
# PATH 2: Use the server-to-server filtered path
if via_server:
# Use the server-to-server filtered path
child_results = await mounted.server._list_tools()
else: # mode == "inventory"
# PATH 1: Use the manager-to-manager unfiltered path
else:
# Use the manager-to-manager unfiltered path
child_results = await mounted.server._tool_manager._list_tools()
# The combination logic is the same for both paths
@ -103,13 +101,13 @@ class ToolManager:
"""
Gets the complete, unfiltered inventory of all tools.
"""
return await self._load_tools(mode="inventory")
return await self._load_tools(via_server=False)
async def _list_tools(self) -> list[Tool]:
"""
Lists all tools, applying protocol filtering.
"""
tools_dict = await self._load_tools(mode="protocol")
tools_dict = await self._load_tools(via_server=True)
return list(tools_dict.values())
def add_tool_from_fn(