113 lines
3.9 KiB
Python
113 lines
3.9 KiB
Python
"""Resource manager functionality."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from pydantic import AnyUrl
|
|
|
|
from mcp.server.mcpserver.resources.base import Resource
|
|
from mcp.server.mcpserver.resources.templates import ResourceTemplate
|
|
from mcp.server.mcpserver.utilities.logging import get_logger
|
|
from mcp.types import Annotations, Icon
|
|
|
|
if TYPE_CHECKING:
|
|
from mcp.server.mcpserver.server import Context
|
|
from mcp.server.session import ServerSessionT
|
|
from mcp.shared.context import LifespanContextT, RequestT
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class ResourceManager:
|
|
"""Manages MCPServer resources."""
|
|
|
|
def __init__(self, warn_on_duplicate_resources: bool = True):
|
|
self._resources: dict[str, Resource] = {}
|
|
self._templates: dict[str, ResourceTemplate] = {}
|
|
self.warn_on_duplicate_resources = warn_on_duplicate_resources
|
|
|
|
def add_resource(self, resource: Resource) -> Resource:
|
|
"""Add a resource to the manager.
|
|
|
|
Args:
|
|
resource: A Resource instance to add
|
|
|
|
Returns:
|
|
The added resource. If a resource with the same URI already exists,
|
|
returns the existing resource.
|
|
"""
|
|
logger.debug(
|
|
"Adding resource",
|
|
extra={
|
|
"uri": resource.uri,
|
|
"type": type(resource).__name__,
|
|
"resource_name": resource.name,
|
|
},
|
|
)
|
|
existing = self._resources.get(str(resource.uri))
|
|
if existing:
|
|
if self.warn_on_duplicate_resources:
|
|
logger.warning(f"Resource already exists: {resource.uri}")
|
|
return existing
|
|
self._resources[str(resource.uri)] = resource
|
|
return resource
|
|
|
|
def add_template(
|
|
self,
|
|
fn: Callable[..., Any],
|
|
uri_template: str,
|
|
name: str | None = None,
|
|
title: str | None = None,
|
|
description: str | None = None,
|
|
mime_type: str | None = None,
|
|
icons: list[Icon] | None = None,
|
|
annotations: Annotations | None = None,
|
|
meta: dict[str, Any] | None = None,
|
|
) -> ResourceTemplate:
|
|
"""Add a template from a function."""
|
|
template = ResourceTemplate.from_function(
|
|
fn,
|
|
uri_template=uri_template,
|
|
name=name,
|
|
title=title,
|
|
description=description,
|
|
mime_type=mime_type,
|
|
icons=icons,
|
|
annotations=annotations,
|
|
meta=meta,
|
|
)
|
|
self._templates[template.uri_template] = template
|
|
return template
|
|
|
|
async def get_resource(
|
|
self, uri: AnyUrl | str, context: Context[ServerSessionT, LifespanContextT, RequestT] | None = None
|
|
) -> Resource:
|
|
"""Get resource by URI, checking concrete resources first, then templates."""
|
|
uri_str = str(uri)
|
|
logger.debug("Getting resource", extra={"uri": uri_str})
|
|
|
|
# First check concrete resources
|
|
if resource := self._resources.get(uri_str):
|
|
return resource
|
|
|
|
# Then check templates
|
|
for template in self._templates.values():
|
|
if params := template.matches(uri_str):
|
|
try:
|
|
return await template.create_resource(uri_str, params, context=context)
|
|
except Exception as e: # pragma: no cover
|
|
raise ValueError(f"Error creating resource from template: {e}")
|
|
|
|
raise ValueError(f"Unknown resource: {uri}")
|
|
|
|
def list_resources(self) -> list[Resource]:
|
|
"""List all registered resources."""
|
|
logger.debug("Listing resources", extra={"count": len(self._resources)})
|
|
return list(self._resources.values())
|
|
|
|
def list_templates(self) -> list[ResourceTemplate]:
|
|
"""List all registered templates."""
|
|
logger.debug("Listing templates", extra={"count": len(self._templates)})
|
|
return list(self._templates.values())
|