refactor to support more flexible tool toggles

Author: ischneider

src/planet_mcp/clients.py

@@ -0,0 +1,20 @@
+from planet import Session, __version__
+
+_session: Session | None = None
+
+
+def _init_session(cls: type[Session]) -> Session:
+    _session = cls()
+    # todo add our version?
+    _session._client.headers["User-Agent"] = f"planet-mcp {__version__}"
+    return _session
+
+
+def session() -> Session:
+    """session returns a global Session. dependencies should only call this
+    in their tool/resource calls or within a callable `mcp` initializer
+    """
+    global _session
+    if _session is None:
+        _session = _init_session(Session)
+    return _session

src/planet_mcp/main.py

@@ -3,18 +3,40 @@
 is installed, this is installed as an executable named planet-mcp.
 """
 
-from fastmcp import FastMCP
-import planet
-from planet_mcp.server import MCPService
-
 # note - the mcp dev tooling (e.g. uv run fastmcp dev src/main.py)
 # wants to find a server object named `mcp` (or it won't work)
-mcp = FastMCP("Planet local/stdio MCP")
-sess = planet.Session()
-svc = MCPService(mcp=mcp, session=sess)
+import argparse
+import asyncio
+from planet_mcp.server import init
+
+
+def parse_args() -> argparse.Namespace:
+    def csv(value):
+        return set(t.strip() for t in (value or "").split(","))
+
+    parser = argparse.ArgumentParser(
+        description="Planet MCP Server",
+    )
+    parser.add_argument("args", nargs="*")
+    parser.add_argument("--include-tags", type=csv, default=None)
+    parser.add_argument("--exclude-tags", type=csv, default=None)
+    parser.add_argument("--servers", type=csv, default=None)
+    return parser.parse_args()
+
+
+args = parse_args()
+mcp = init(
+    enabled_servers=args.servers,
+    include_tags=args.include_tags,
+    exclude_tags=args.exclude_tags,
+)
 
 
 # this is the entry point for the executable script installed via package
 # and also supports execution via `uv run fastmcp run src/main.py`
 def main():
-    svc.mcp.run(transport="stdio")
+    mcp.run(transport="stdio")
+
+
+if __name__ == "__main__":
+    main()

src/planet_mcp/server.py

@@ -1,142 +1,57 @@
-import functools
-import inspect
-from planet_mcp import descriptions, tiles
+import asyncio
+from contextlib import asynccontextmanager
+from typing import AsyncIterator
 from fastmcp import FastMCP
-import planet
-from typing import Union, Optional
+from planet_mcp import servers
 
-from pydantic import PydanticSchemaGenerationError
+_instructions = """
+Instructions for using Planet's official MCP server.
+"""
 
 
-DEFAULT_IGNORE = {
-    "data_wait_asset",
-    "orders_wait",
-}
+class PlanetContext:
+    def __init__(self):
+        pass
 
 
-class MCPService:
-    """MCP service for the Planet API
+def _lifespan(context: PlanetContext):
 
-    member `mcp` contains the FastMCP app instance.
-    """
+    @asynccontextmanager
+    async def lifespan(server: FastMCP) -> AsyncIterator[PlanetContext]:
+        yield context
 
-    def __init__(
-        self,
-        mcp: FastMCP,
-        session: planet.Session,
-        ignore: set[str] | None = None,
-        include: set[str] | None = None,
-    ):
-        self.mcp = mcp
-        self.session = session
+    return lifespan
 
-        if ignore and include:
-            raise ValueError("Cannot specify both ignore and include sets.")
 
-        self.include = include
-        self.ignore = ignore
+def init(
+    enabled_servers: set[str] | None = None,
+    include_tags: set[str] | None = None,
+    exclude_tags: set[str] | None = None,
+) -> FastMCP:
 
-        if self.ignore is None and self.include is None:
-            self.ignore = DEFAULT_IGNORE
-        self.ignore = ignore if ignore is not None else DEFAULT_IGNORE
+    context = PlanetContext()
 
-        self.make_tools(planet.FeaturesClient, "features")
-        self.make_tools(planet.DataClient, "data")
-        self.make_tools(planet.OrdersClient, "orders")
-        self.make_tools(planet.SubscriptionsClient, "subscriptions")
-        self.make_tools(planet.MosaicsClient, "mosaics")
-        self.make_tools(planet.DestinationsClient, "destinations")
+    mcp = FastMCP(
+        "Planet MCP Server",
+        lifespan=_lifespan(context),
+        instructions=_instructions,
+        include_tags=include_tags,
+        exclude_tags=exclude_tags,
+    )
 
-        # add tiles tools
-        tiles.init(self.mcp, self.session)
+    for server in servers.all:
+        try:
+            # server protocol is either a variable or callable named mcp
+            entry = getattr(server, "mcp")
+        except AttributeError:
+            raise Exception(f"programmer error, mcp attribute not in {server}")
+        if callable(entry):
+            entry = entry()
+        if not isinstance(entry, FastMCP):
+            raise Exception(
+                f"programmer error, expected FastMCP type, got {type(entry)}"
+            )
+        if enabled_servers is None or entry.name in enabled_servers:
+            mcp.mount(entry, entry.name)  # type: ignore
 
-    def make_tools(self, client_class, prefix: str):
-        for name, func in inspect.getmembers(client_class(self.session)):
-            if inspect.ismethod(func) and name[0] != "_":
-                full_name = f"{prefix}_{name}"
-
-                if self.ignore is not None and full_name in self.ignore:
-                    continue
-                if self.include is not None and name not in self.include:
-                    continue
-
-                # extended tool options
-                opts = {}
-
-                # check if there is a description override for this tool
-                if full_name in descriptions.overrides:
-                    opts["description"] = descriptions.overrides[full_name]
-
-                # async generator functions have an incompatible return type.
-                # ensure they are converted to a list[dict] return type.
-                if inspect.isasyncgenfunction(func):
-                    func = _async_get_wrapper(func, prefix)
-
-                # no return functions end up with a "self" parameter so this
-                # works around by adding a simple response
-                # @todo - upstream bug?
-                sig = inspect.signature(func)
-                if sig.return_annotation is None:
-                    func = _return_wrapper(func)
-
-                try:
-                    self.mcp.tool(func, name=full_name, **opts)
-                except PydanticSchemaGenerationError:
-                    # there's a few functions we need to modify again because of custom types.
-                    modified_func = _create_param_modified_wrapper(func)
-                    try:
-                        self.mcp.tool(modified_func, name=full_name, **opts)
-                    except Exception as e:
-                        print("Unable to add tool", full_name, e)
-
-
-def _async_get_wrapper(f, prefix):
-    """wrap an async generator to return a list[dict]"""
-
-    @functools.wraps(f)
-    async def generate_async(*args, **kwargs) -> list[dict]:
-        return [i async for i in (f(*args, **kwargs))]
-
-    # functool.wraps annotates using the original function return
-    generate_async.__annotations__["return"] = list[dict]
-    return generate_async
-
-
-def _return_wrapper(func):
-    @functools.wraps(func)
-    async def wrapper(*args, **kwargs):
-        await func(*args, **kwargs)
-        return {"status": "ok"}
-
-    wrapper.__annotations__["return"] = dict
-    return wrapper
-
-
-def _create_param_modified_wrapper(original_func):
-    """
-    Some functions that accept special types (typing.Protocol) fail during
-    FastMCP tool registration. This wrapper modifies the function's signature,
-    replacing the type hints with simple types that FastMCP can handle.
-    """
-
-    @functools.wraps(original_func)
-    async def wrapper(*args, **kwargs):
-        return await original_func(*args, **kwargs)
-
-    try:
-        sig = inspect.signature(original_func)
-
-        for param_name, param in sig.parameters.items():
-            # Convert problematic types into useable ones.
-            # we want to override the GeojsonLike field
-            # and remove the union of the Feature which accepts a dict or a GeoInterface
-            if param_name in ("feature", "quad", "mosaic", "series"):
-                wrapper.__annotations__[param_name] = dict
-            elif param_name == "geometry" and "planet.models" in str(param.annotation):
-                wrapper.__annotations__[param_name] = Optional[Union[dict, str]] | None
-
-    except Exception as e:
-        print(f"Error modifying signature: {e}")
-        wrapper.__annotations__ = {}
-
-    return wrapper
+    return mcp

src/planet_mcp/servers/__init__.py

@@ -0,0 +1,7 @@
+from . import sdk_wrapper
+from . import tiles
+
+all = [
+    sdk_wrapper,
+    tiles,
+]

src/planet_mcp/descriptions.py renamed to src/planet_mcp/servers/descriptions.py

@@ -0,0 +1,120 @@
+import functools
+import inspect
+
+from planet_mcp.clients import session
+from . import descriptions
+from fastmcp import FastMCP
+import planet
+from typing import Union, Optional
+
+from pydantic import PydanticSchemaGenerationError
+
+
+_DEFAULT_IGNORE = {
+    "data_wait_asset",
+    "orders_wait",
+}
+
+
+def mcp() -> FastMCP:
+    mcp = FastMCP("sdk_wrapper")
+    make_tools(mcp, planet.FeaturesClient, "features")
+    make_tools(mcp, planet.DataClient, "data")
+    make_tools(mcp, planet.OrdersClient, "orders")
+    make_tools(mcp, planet.SubscriptionsClient, "subscriptions")
+    make_tools(mcp, planet.MosaicsClient, "mosaics")
+    make_tools(mcp, planet.DestinationsClient, "destinations")
+    return mcp
+
+
+def make_tools(mcp: FastMCP, client_class: type, prefix: str):
+    for name, func in inspect.getmembers(client_class(session())):
+        if inspect.ismethod(func) and name[0] != "_":
+            full_name = f"{prefix}_{name}"
+
+            if full_name in _DEFAULT_IGNORE:
+                continue
+
+            # extended tool options
+            opts = {}
+
+            # check if there is a description override for this tool
+            if full_name in descriptions.overrides:
+                opts["description"] = descriptions.overrides[full_name]
+
+            # async generator functions have an incompatible return type.
+            # ensure they are converted to a list[dict] return type.
+            if inspect.isasyncgenfunction(func):
+                func = _async_get_wrapper(func, prefix)
+
+            # no return functions end up with a "self" parameter so this
+            # works around by adding a simple response
+            # @todo - upstream bug?
+            sig = inspect.signature(func)
+            if sig.return_annotation is None:
+                func = _return_wrapper(func)
+
+            if "download" in name:
+                opts["tags"] = set("download")
+
+            try:
+                mcp.tool(func, name=full_name, **opts)
+            except PydanticSchemaGenerationError:
+                # there's a few functions we need to modify again because of custom types.
+                modified_func = _create_param_modified_wrapper(func)
+                try:
+                    mcp.tool(modified_func, name=full_name, **opts)
+                except Exception as e:
+                    print("Unable to add tool", full_name, e)
+
+
+def _async_get_wrapper(f, prefix):
+    """wrap an async generator to return a list[dict]"""
+
+    @functools.wraps(f)
+    async def generate_async(*args, **kwargs) -> list[dict]:
+        return [i async for i in (f(*args, **kwargs))]
+
+    # functool.wraps annotates using the original function return
+    generate_async.__annotations__["return"] = list[dict]
+    return generate_async
+
+
+def _return_wrapper(func):
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        await func(*args, **kwargs)
+        return {"status": "ok"}
+
+    wrapper.__annotations__["return"] = dict
+    return wrapper
+
+
+def _create_param_modified_wrapper(original_func):
+    """
+    Some functions that accept special types (typing.Protocol) fail during
+    FastMCP tool registration. This wrapper modifies the function's signature,
+    replacing the type hints with simple types that FastMCP can handle.
+    """
+
+    @functools.wraps(original_func)
+    async def wrapper(*args, **kwargs):
+        return await original_func(*args, **kwargs)
+
+    try:
+        sig = inspect.signature(original_func)
+
+        for param_name, param in sig.parameters.items():
+            # Convert problematic types into useable ones.
+            # we want to override the GeojsonLike field
+            # and remove the union of the Feature which accepts a dict or a GeoInterface
+            if param_name in ("feature", "quad", "mosaic", "series"):
+                wrapper.__annotations__[param_name] = dict
+            elif param_name == "geometry" and "planet.models" in str(param.annotation):
+                wrapper.__annotations__[param_name] = Optional[Union[dict, str]] | None
+
+    except Exception as e:
+        print(f"Error modifying signature: {e}")
+        wrapper.__annotations__ = {}
+
+    return wrapper