Merge pull request #10 from jkawamoto/fix-tools

Refactor server setup and ignore type error

Author: jkawamoto

src/mcp_bear/server.py

@@ -13,21 +13,23 @@
 from contextlib import asynccontextmanager
 from copy import deepcopy
 from dataclasses import dataclass
+from functools import partial
 from http import HTTPStatus
 from typing import cast, AsyncIterator
 from urllib.parse import urlencode, quote, unquote_plus
 
 from fastapi import FastAPI, Request
 from mcp.server import FastMCP
 from mcp.server.fastmcp import Context
-from mcp.server.session import ServerSessionT
 from pydantic import Field
 from starlette.datastructures import QueryParams
 from uvicorn import Config, Server
 from uvicorn.config import LOGGING_CONFIG
 
 BASE_URL = "bear://x-callback-url"
 
+LOGGER = logging.getLogger(__name__)
+
 
 @dataclass
 class ErrorResponse(Exception):
@@ -79,53 +81,52 @@ class AppContext:
     grab_url_results: Queue[Future[QueryParams]]
 
 
-def create_server(token: str, callback_host: str, callback_port: int) -> FastMCP:
-    logger = logging.getLogger(__name__)
-
-    @asynccontextmanager
-    async def app_lifespan(_server: FastMCP) -> AsyncIterator[AppContext]:
-        callback = FastAPI()
-
-        log_config = deepcopy(LOGGING_CONFIG)
-        log_config["handlers"]["access"]["stream"] = "ext://sys.stderr"
-        server = Server(
-            Config(
-                app=callback,
-                host=callback_host,
-                port=callback_port,
-                log_level="warning",
-                log_config=log_config,
-            )
+@asynccontextmanager
+async def app_lifespan(_server: FastMCP, callback_host: str, callback_port: int) -> AsyncIterator[AppContext]:
+    callback = FastAPI()
+
+    log_config = deepcopy(LOGGING_CONFIG)
+    log_config["handlers"]["access"]["stream"] = "ext://sys.stderr"
+    server = Server(
+        Config(
+            app=callback,
+            host=callback_host,
+            port=callback_port,
+            log_level="warning",
+            log_config=log_config,
         )
+    )
+
+    LOGGER.info(f"Starting callback server on {callback_host}:{callback_port}")
+    server_task = asyncio.create_task(server.serve())
+    try:
+        yield AppContext(
+            open_note_results=register_callback(callback, "open-note"),
+            create_results=register_callback(callback, "create"),
+            tags_results=register_callback(callback, "tags"),
+            open_tag_results=register_callback(callback, "open-tag"),
+            todo_results=register_callback(callback, "todo"),
+            today_results=register_callback(callback, "today"),
+            search_results=register_callback(callback, "search"),
+            grab_url_results=register_callback(callback, "grab-url"),
+        )
+    finally:
+        LOGGER.info("Stopping callback server")
+        server.should_exit = True
+        await server_task
 
-        logger.info(f"Starting callback server on {callback_host}:{callback_port}")
-        server_task = asyncio.create_task(server.serve())
-        try:
-            yield AppContext(
-                open_note_results=register_callback(callback, "open-note"),
-                create_results=register_callback(callback, "create"),
-                tags_results=register_callback(callback, "tags"),
-                open_tag_results=register_callback(callback, "open-tag"),
-                todo_results=register_callback(callback, "todo"),
-                today_results=register_callback(callback, "today"),
-                search_results=register_callback(callback, "search"),
-                grab_url_results=register_callback(callback, "grab-url"),
-            )
-        finally:
-            logger.info("Stopping callback server")
-            server.should_exit = True
-            await server_task
 
-    mcp = FastMCP("Bear", lifespan=app_lifespan)
+def create_server(token: str, callback_host: str, callback_port: int) -> FastMCP:
+    mcp = FastMCP("Bear", lifespan=partial(app_lifespan, callback_host=callback_host, callback_port=callback_port))
 
     @mcp.tool()
     async def open_note(
-        ctx: Context[ServerSessionT, AppContext],
+        ctx: Context,
         id: str | None = Field(description="note unique identifier", default=None),
         title: str | None = Field(description="note title", default=None),
     ) -> str:
         """Open a note identified by its title or id and return its content."""
-        app_ctx = ctx.request_context.lifespan_context
+        app_ctx: AppContext = ctx.request_context.lifespan_context  # type: ignore
         future = Future[QueryParams]()
         await app_ctx.open_note_results.put(future)
 
@@ -152,14 +153,14 @@ async def open_note(
 
     @mcp.tool()
     async def create(
-        ctx: Context[ServerSessionT, AppContext],
+        ctx: Context,
         title: str | None = Field(description="note title", default=None),
         text: str | None = Field(description="note body", default=None),
         tags: list[str] | None = Field(description="list of tags", default=None),
         timestamp: bool = Field(description="prepend the current date and time to the text", default=False),
     ) -> str:
         """Create a new note and return its unique identifier. Empty notes are not allowed."""
-        app_ctx = ctx.request_context.lifespan_context
+        app_ctx: AppContext = ctx.request_context.lifespan_context  # type: ignore
         future = Future[QueryParams]()
         await app_ctx.create_results.put(future)
 
@@ -187,10 +188,10 @@ async def create(
 
     @mcp.tool()
     async def tags(
-        ctx: Context[ServerSessionT, AppContext],
+        ctx: Context,
     ) -> list[str]:
         """Return all the tags currently displayed in Bear’s sidebar."""
-        app_ctx = ctx.request_context.lifespan_context
+        app_ctx: AppContext = ctx.request_context.lifespan_context  # type: ignore
         future = Future[QueryParams]()
         await app_ctx.tags_results.put(future)
 
@@ -208,11 +209,11 @@ async def tags(
 
     @mcp.tool()
     async def open_tag(
-        ctx: Context[ServerSessionT, AppContext],
+        ctx: Context,
         name: str = Field(description="tag name or a list of tags divided by comma"),
     ) -> list[str]:
         """Show all the notes which have a selected tag in bear."""
-        app_ctx = ctx.request_context.lifespan_context
+        app_ctx: AppContext = ctx.request_context.lifespan_context  # type: ignore
         future = Future[QueryParams]()
         await app_ctx.open_tag_results.put(future)
 
@@ -231,11 +232,11 @@ async def open_tag(
 
     @mcp.tool()
     async def todo(
-        ctx: Context[ServerSessionT, AppContext],
+        ctx: Context,
         search: str | None = Field(description="string to search", default=None),
     ) -> list[str]:
         """Select the Todo sidebar item."""
-        app_ctx = ctx.request_context.lifespan_context
+        app_ctx: AppContext = ctx.request_context.lifespan_context  # type: ignore
         future = Future[QueryParams]()
         await app_ctx.todo_results.put(future)
 
@@ -256,11 +257,11 @@ async def todo(
 
     @mcp.tool()
     async def today(
-        ctx: Context[ServerSessionT, AppContext],
+        ctx: Context,
         search: str | None = Field(description="string to search", default=None),
     ) -> list[str]:
         """Select the Today sidebar item."""
-        app_ctx = ctx.request_context.lifespan_context
+        app_ctx: AppContext = ctx.request_context.lifespan_context  # type: ignore
         future = Future[QueryParams]()
         await app_ctx.today_results.put(future)
 
@@ -281,12 +282,12 @@ async def today(
 
     @mcp.tool()
     async def search(
-        ctx: Context[ServerSessionT, AppContext],
+        ctx: Context,
         term: str | None = Field(description="string to search", default=None),
         tag: str | None = Field(description="tag to search into", default=None),
     ) -> list[str]:
         """Show search results in Bear for all notes or for a specific tag."""
-        app_ctx = ctx.request_context.lifespan_context
+        app_ctx: AppContext = ctx.request_context.lifespan_context  # type: ignore
         future = Future[QueryParams]()
         await app_ctx.search_results.put(future)
 
@@ -309,15 +310,15 @@ async def search(
 
     @mcp.tool()
     async def grab_url(
-        ctx: Context[ServerSessionT, AppContext],
+        ctx: Context,
         url: str = Field(description="url to grab"),
         tags: list[str] | None = Field(
             description="list of tags. If tags are specified in the Bear’s web content preferences, this parameter is ignored.",
             default=None,
         ),
     ) -> str:
         """Create a new note with the content of a web page and return its unique identifier."""
-        app_ctx = ctx.request_context.lifespan_context
+        app_ctx: AppContext = ctx.request_context.lifespan_context  # type: ignore
         future = Future[QueryParams]()
         await app_ctx.grab_url_results.put(future)