fix(review): harden tool validation and lifecycle cleanup

Author: ayhammouda

src/mcp_server_python_docs/__main__.py

@@ -9,13 +9,40 @@
 # === STDIO HYGIENE (HYGN-01, B3 blocker) ===
 # These MUST be the first imports and operations, before anything
 # that might write to stdout.
+import atexit
 import os
 import signal
 import sys
 
 # Save the real stdout fd for the MCP framer, then redirect fd 1 to stderr.
 # After this, any print() or write to fd 1 goes to stderr, not the MCP pipe.
-_real_stdout_fd = os.dup(1)
+_saved_stdout_fd: int | None = os.dup(1)
+
+
+def _close_saved_stdout_fd() -> None:
+    """Close the saved stdout fd when the CLI exits without serving."""
+    global _saved_stdout_fd
+    if _saved_stdout_fd is None:
+        return
+    try:
+        os.close(_saved_stdout_fd)
+    except OSError:
+        pass
+    finally:
+        _saved_stdout_fd = None
+
+
+def _consume_saved_stdout_fd() -> int:
+    """Hand off the saved stdout fd to the stdio MCP transport."""
+    global _saved_stdout_fd
+    if _saved_stdout_fd is None:
+        raise RuntimeError("Saved stdout fd is not available")
+    fd = _saved_stdout_fd
+    _saved_stdout_fd = None
+    return fd
+
+
+atexit.register(_close_saved_stdout_fd)
 os.dup2(2, 1)
 sys.stdout = sys.stderr
 
@@ -59,12 +86,13 @@ def serve() -> None:
     from mcp_server_python_docs.server import create_server
 
     mcp_server = create_server()
+    saved_stdout_fd = _consume_saved_stdout_fd()
 
     # Restore the real stdout fd for MCP protocol framing.
     # By this point all imports are done — no third-party code will
     # print to stdout during MCP communication.
-    os.dup2(_real_stdout_fd, 1)
-    os.close(_real_stdout_fd)
+    os.dup2(saved_stdout_fd, 1)
+    os.close(saved_stdout_fd)
 
     try:
         mcp_server.run(transport="stdio")

src/mcp_server_python_docs/ingestion/publish.py

@@ -26,15 +26,15 @@
 def generate_build_path() -> Path:
     """Generate a timestamped build artifact path (PUBL-01).
 
-    Returns a path like ``~/.cache/mcp-python-docs/build-20260416-143022.db``.
+    Returns a path like ``~/.cache/mcp-python-docs/build-20260416-143022-123456.db``.
     Creates the cache directory if it does not exist.
 
     Returns:
         Path to the new build artifact.
     """
     cache_dir = get_cache_dir()
     cache_dir.mkdir(parents=True, exist_ok=True)
-    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
+    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S-%f")
     return cache_dir / f"build-{timestamp}.db"
 
 

src/mcp_server_python_docs/retrieval/budget.py

@@ -31,7 +31,7 @@ def apply_budget(
         next_start_index is None when not truncated (all remaining text fits).
     """
     if not text or max_chars <= 0:
-        return ("", bool(text), 0 if text else None)
+        return ("", False, None)
 
     if start_index >= len(text):
         return ("", False, None)

src/mcp_server_python_docs/retrieval/query.py

@@ -174,8 +174,10 @@ def build_match_expression(
 
     # If expansion added new terms, OR-join all terms
     if expanded != original_tokens:
+        original_query = fts5_escape(query)
         escaped_terms = [fts5_escape(term) for term in sorted(expanded)]
-        return " OR ".join(escaped_terms)
+        extra_terms = [term for term in escaped_terms if term != original_query]
+        return " OR ".join([original_query, *extra_terms])
 
     # No expansion -- use plain escaped query (implicit AND)
     return fts5_escape(query)

src/mcp_server_python_docs/server.py

@@ -13,13 +13,14 @@
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import Literal
+from typing import Annotated, Literal
 
 import platformdirs
 import yaml
 from mcp.server.fastmcp import Context, FastMCP
 from mcp.server.fastmcp.exceptions import ToolError
 from mcp.types import ToolAnnotations
+from pydantic import Field
 
 from mcp_server_python_docs.app_context import AppContext
 from mcp_server_python_docs.detection import detect_python_version, match_to_indexed
@@ -81,50 +82,51 @@ async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
     # Open read-only connection (STOR-06, STOR-07)
     db = get_readonly_connection(index_path)
 
-    # Check FTS5 (STOR-08)
-    _assert_fts5(db)
-
-    # Construct service instances (Phase 5 — service layer wiring)
-    search_svc = SearchService(db, synonyms)
-    content_svc = ContentService(db)
-    version_svc = VersionService(db)
-
-    # Detect user's Python version and match to indexed versions
-    detected_ver, detected_src = detect_python_version()
-    indexed_versions = [
-        r[0] for r in db.execute("SELECT version FROM doc_sets ORDER BY version").fetchall()
-    ]
-    matched = match_to_indexed(detected_ver, indexed_versions)
-    if matched:
-        logger.info("User Python %s matches indexed version — using as default", matched)
-    else:
-        logger.info(
-            "User Python %s not in index %s — using normal default",
-            detected_ver,
-            indexed_versions,
-        )
-
     try:
-        yield AppContext(
-            db=db,
-            index_path=index_path,
-            synonyms=synonyms,
-            search_service=search_svc,
-            content_service=content_svc,
-            version_service=version_svc,
-            detected_python_version=matched,
-            detected_python_source=detected_src,
-        )
-    except Exception:
-        # HYGN-05: log lifespan errors, write last-error.log, re-raise original
-        error_msg = traceback.format_exc()
-        logger.error("Lifespan error: %s", error_msg)
+        # Check FTS5 (STOR-08)
+        _assert_fts5(db)
+
+        # Construct service instances (Phase 5 — service layer wiring)
+        search_svc = SearchService(db, synonyms)
+        content_svc = ContentService(db)
+        version_svc = VersionService(db)
+
+        # Detect user's Python version and match to indexed versions
+        detected_ver, detected_src = detect_python_version()
+        indexed_versions = [
+            r[0] for r in db.execute("SELECT version FROM doc_sets ORDER BY version").fetchall()
+        ]
+        matched = match_to_indexed(detected_ver, indexed_versions)
+        if matched:
+            logger.info("User Python %s matches indexed version — using as default", matched)
+        else:
+            logger.info(
+                "User Python %s not in index %s — using normal default",
+                detected_ver,
+                indexed_versions,
+            )
+
         try:
-            error_log = cache_dir / "last-error.log"
-            error_log.write_text(error_msg)
+            yield AppContext(
+                db=db,
+                index_path=index_path,
+                synonyms=synonyms,
+                search_service=search_svc,
+                content_service=content_svc,
+                version_service=version_svc,
+                detected_python_version=matched,
+                detected_python_source=detected_src,
+            )
         except Exception:
-            pass
-        raise
+            # HYGN-05: log lifespan errors, write last-error.log, re-raise original
+            error_msg = traceback.format_exc()
+            logger.error("Lifespan error: %s", error_msg)
+            try:
+                error_log = cache_dir / "last-error.log"
+                error_log.write_text(error_msg)
+            except Exception:
+                pass
+            raise
     finally:
         db.close()
 
@@ -137,6 +139,47 @@ async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
     openWorldHint=False,
 )
 
+SearchQueryParam = Annotated[
+    str,
+    Field(
+        max_length=500,
+        description="Search query - Python symbol (asyncio.TaskGroup) or concept (parse json)",
+    ),
+]
+VersionParam = Annotated[
+    str | None,
+    Field(description="Python version (e.g. '3.13'). Defaults to latest."),
+]
+SearchKindParam = Annotated[
+    Literal["auto", "page", "symbol", "section", "example"],
+    Field(
+        description=(
+            "Search type. Use 'symbol' for API lookups, "
+            "'example' for code samples, 'auto' otherwise."
+        )
+    ),
+]
+MaxResultsParam = Annotated[
+    int,
+    Field(ge=1, le=20, description="Maximum number of results to return."),
+]
+SlugParam = Annotated[
+    str,
+    Field(max_length=500, description="Page slug (e.g. 'library/asyncio-task.html')"),
+]
+AnchorParam = Annotated[
+    str | None,
+    Field(description="Section anchor for section-only retrieval"),
+]
+MaxCharsParam = Annotated[
+    int,
+    Field(ge=100, le=50000, description="Maximum characters to return"),
+]
+StartIndexParam = Annotated[
+    int,
+    Field(ge=0, description="Start position for pagination"),
+]
+
 
 def create_server() -> FastMCP:
     """Create and configure the FastMCP server."""
@@ -147,10 +190,10 @@ def create_server() -> FastMCP:
 
     @mcp.tool(annotations=_TOOL_ANNOTATIONS)
     def search_docs(
-        query: str,
-        version: str | None = None,
-        kind: Literal["auto", "page", "symbol", "section", "example"] = "auto",
-        max_results: int = 5,
+        query: SearchQueryParam,
+        version: VersionParam = None,
+        kind: SearchKindParam = "auto",
+        max_results: MaxResultsParam = 5,
         ctx: Context = None,  # type: ignore[assignment]
     ) -> SearchDocsResult:
         """Search Python documentation. Use kind='symbol' for API lookups
@@ -168,11 +211,11 @@ def search_docs(
 
     @mcp.tool(annotations=_TOOL_ANNOTATIONS)
     def get_docs(
-        slug: str,
-        version: str | None = None,
-        anchor: str | None = None,
-        max_chars: int = 8000,
-        start_index: int = 0,
+        slug: SlugParam,
+        version: VersionParam = None,
+        anchor: AnchorParam = None,
+        max_chars: MaxCharsParam = 8000,
+        start_index: StartIndexParam = 0,
         ctx: Context = None,  # type: ignore[assignment]
     ) -> GetDocsResult:
         """Retrieve a documentation page or specific section. Provide anchor for
@@ -214,7 +257,6 @@ def detect_python_version(
         matches an indexed documentation set."""
         app_ctx: AppContext = ctx.request_context.lifespan_context
         detected_ver = app_ctx.detected_python_version
-        detected_src = app_ctx.detected_python_source or "unknown"
 
         # Re-run detection to get the raw version even if it didn't match
         from mcp_server_python_docs.detection import detect_python_version as _detect

src/mcp_server_python_docs/services/observability.py

@@ -5,7 +5,8 @@
 unstable, so we instrument at the service layer.
 
 Log format (logfmt — D-10 from Phase 1):
-  tool=search_docs version=3.13 latency_ms=12 result_count=5 truncated=false resolution=fts synonym_expansion=yes
+  tool=search_docs version=3.13 latency_ms=12 result_count=5
+  truncated=false resolution=fts synonym_expansion=yes
 """
 from __future__ import annotations
 
@@ -14,6 +15,7 @@
 import sys
 import time
 from collections.abc import Callable
+from types import TracebackType
 from typing import Any
 
 
@@ -56,8 +58,15 @@ def decorator(fn: Callable) -> Callable:
         @functools.wraps(fn)
         def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
             start = time.monotonic()
+            result: Any = None
+            error: Exception | None = None
+            error_tb: TracebackType | None = None
 
-            result = fn(self, *args, **kwargs)
+            try:
+                result = fn(self, *args, **kwargs)
+            except Exception as exc:
+                error = exc
+                error_tb = exc.__traceback__
 
             elapsed_ms = (time.monotonic() - start) * 1000
 
@@ -77,37 +86,43 @@ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
                 version_val = kwargs.get("version")
             fields["version"] = version_val or "default"
 
-            # Extract result-specific fields
-            if hasattr(result, "hits"):
-                # SearchDocsResult
-                fields["result_count"] = len(result.hits)
-                # Resolution path from service state
-                if hasattr(self, "_last_resolution"):
-                    fields["resolution"] = self._last_resolution
-                else:
-                    fields["resolution"] = "fts"
-                fields["truncated"] = False
-            elif hasattr(result, "truncated"):
-                # GetDocsResult
-                fields["result_count"] = 1 if result.content else 0
-                fields["truncated"] = result.truncated
-                fields["resolution"] = "exact"
-            elif hasattr(result, "versions"):
-                # ListVersionsResult
-                fields["result_count"] = len(result.versions)
-                fields["truncated"] = False
-                fields["resolution"] = "exact"
-
-            # Synonym expansion detection from service state
-            if hasattr(self, "_last_synonym_expanded"):
-                fields["synonym_expansion"] = (
-                    "yes" if self._last_synonym_expanded else "no"
-                )
+            if error is None:
+                # Extract result-specific fields
+                if hasattr(result, "hits"):
+                    # SearchDocsResult
+                    fields["result_count"] = len(result.hits)
+                    # Resolution path from service state
+                    if hasattr(self, "_last_resolution"):
+                        fields["resolution"] = self._last_resolution
+                    else:
+                        fields["resolution"] = "fts"
+                    fields["truncated"] = False
+                elif hasattr(result, "truncated"):
+                    # GetDocsResult
+                    fields["result_count"] = 1 if result.content else 0
+                    fields["truncated"] = result.truncated
+                    fields["resolution"] = "exact"
+                elif hasattr(result, "versions"):
+                    # ListVersionsResult
+                    fields["result_count"] = len(result.versions)
+                    fields["truncated"] = False
+                    fields["resolution"] = "exact"
+
+                # Synonym expansion detection from service state
+                if hasattr(self, "_last_synonym_expanded"):
+                    fields["synonym_expansion"] = (
+                        "yes" if self._last_synonym_expanded else "no"
+                    )
+            else:
+                fields["error"] = type(error).__name__
 
             # Write logfmt line to stderr (HYGN-01 safe — stderr only)
             log_line = _format_logfmt(**fields)
             print(log_line, file=sys.stderr)
 
+            if error is not None:
+                raise error.with_traceback(error_tb)
+
             return result
 
         return wrapper