syndarix/mcp-servers/knowledge-base/server.py

"""
Knowledge Base MCP Server.

Provides RAG capabilities with pgvector for semantic search,
intelligent chunking, and collection management.
"""

import inspect
import logging
import re
from contextlib import asynccontextmanager
from typing import Any, get_type_hints

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from fastmcp import FastMCP
from pydantic import Field

from collection_manager import CollectionManager, get_collection_manager
from collections.abc import AsyncIterator
from config import get_settings
from database import DatabaseManager, get_database_manager
from embeddings import EmbeddingGenerator, get_embedding_generator
from exceptions import ErrorCode, KnowledgeBaseError
from models import (
    ChunkType,
    DeleteRequest,
    FileType,
    IngestRequest,
    SearchRequest,
    SearchType,
)
from search import SearchEngine, get_search_engine

# Input validation patterns
# Allow alphanumeric, hyphens, underscores (1-128 chars)
ID_PATTERN = re.compile(r"^[a-zA-Z0-9_-]{1,128}$")
# Collection names: alphanumeric, hyphens, underscores (1-64 chars)
COLLECTION_PATTERN = re.compile(r"^[a-zA-Z0-9_-]{1,64}$")


def _validate_id(value: str, field_name: str) -> str | None:
    """Validate project_id or agent_id format.

    Returns error message if invalid, None if valid.
    """
    # Handle FieldInfo objects from direct .fn() calls in tests
    if not isinstance(value, str):
        return f"{field_name} must be a string"
    if not value:
        return f"{field_name} is required"
    if not ID_PATTERN.match(value):
        return f"Invalid {field_name}: must be 1-128 alphanumeric characters, hyphens, or underscores"
    return None


def _validate_collection(value: str) -> str | None:
    """Validate collection name format.

    Returns error message if invalid, None if valid.
    """
    # Handle FieldInfo objects from direct .fn() calls in tests
    if not isinstance(value, str):
        return None  # Non-string means default not resolved, skip validation
    if not COLLECTION_PATTERN.match(value):
        return "Invalid collection: must be 1-64 alphanumeric characters, hyphens, or underscores"
    return None


def _validate_source_path(value: str | None) -> str | None:
    """Validate source_path to prevent path traversal.

    Returns error message if invalid, None if valid.
    """
    if value is None:
        return None

    # Handle FieldInfo objects from direct .fn() calls in tests
    if not isinstance(value, str):
        return None  # Non-string means default not resolved, skip validation

    # Normalize path and check for traversal attempts
    if ".." in value:
        return "Invalid source_path: path traversal not allowed"

    # Check for null bytes (used in some injection attacks)
    if "\x00" in value:
        return "Invalid source_path: null bytes not allowed"

    # Limit path length to prevent DoS
    if len(value) > 4096:
        return "Invalid source_path: path too long (max 4096 chars)"

    return None

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

# Global instances
_database: DatabaseManager | None = None
_embeddings: EmbeddingGenerator | None = None
_search: SearchEngine | None = None
_collections: CollectionManager | None = None


@asynccontextmanager
async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
    """Application lifespan handler."""
    global _database, _embeddings, _search, _collections

    logger.info("Starting Knowledge Base MCP Server...")

    # Initialize database
    _database = get_database_manager()
    await _database.initialize()

    # Initialize embedding generator
    _embeddings = get_embedding_generator()
    await _embeddings.initialize()

    # Initialize search engine
    _search = get_search_engine()

    # Initialize collection manager
    _collections = get_collection_manager()

    logger.info("Knowledge Base MCP Server started successfully")

    yield

    # Cleanup
    logger.info("Shutting down Knowledge Base MCP Server...")

    if _embeddings:
        await _embeddings.close()

    if _database:
        await _database.close()

    logger.info("Knowledge Base MCP Server shut down")


# Create FastMCP server
mcp = FastMCP("syndarix-knowledge-base")

# Create FastAPI app with lifespan
app = FastAPI(
    title="Knowledge Base MCP Server",
    description="RAG with pgvector for semantic search",
    version="0.1.0",
    lifespan=lifespan,
)


@app.get("/health")
async def health_check() -> dict[str, Any]:
    """Health check endpoint.

    Checks all dependencies: database, Redis cache, and LLM Gateway.
    Returns degraded status if any non-critical dependency fails.
    Returns unhealthy status if critical dependencies fail.
    """
    from datetime import UTC, datetime

    status: dict[str, Any] = {
        "status": "healthy",
        "service": "knowledge-base",
        "version": "0.1.0",
        "timestamp": datetime.now(UTC).isoformat(),
        "dependencies": {},
    }

    is_degraded = False
    is_unhealthy = False

    # Check database connection (critical)
    try:
        if _database and _database._pool:
            async with _database.acquire() as conn:
                await conn.fetchval("SELECT 1")
            status["dependencies"]["database"] = "connected"
        else:
            status["dependencies"]["database"] = "not initialized"
            is_unhealthy = True
    except Exception as e:
        status["dependencies"]["database"] = f"error: {e}"
        is_unhealthy = True

    # Check Redis cache (non-critical - degraded without it)
    try:
        if _embeddings and _embeddings._redis:
            await _embeddings._redis.ping()
            status["dependencies"]["redis"] = "connected"
        else:
            status["dependencies"]["redis"] = "not initialized"
            is_degraded = True
    except Exception as e:
        status["dependencies"]["redis"] = f"error: {e}"
        is_degraded = True

    # Check LLM Gateway connectivity (non-critical for health check)
    try:
        if _embeddings and _embeddings._http_client:
            settings = get_settings()
            response = await _embeddings._http_client.get(
                f"{settings.llm_gateway_url}/health",
                timeout=5.0,
            )
            if response.status_code == 200:
                status["dependencies"]["llm_gateway"] = "connected"
            else:
                status["dependencies"]["llm_gateway"] = f"unhealthy (status {response.status_code})"
                is_degraded = True
        else:
            status["dependencies"]["llm_gateway"] = "not initialized"
            is_degraded = True
    except Exception as e:
        status["dependencies"]["llm_gateway"] = f"error: {e}"
        is_degraded = True

    # Set overall status
    if is_unhealthy:
        status["status"] = "unhealthy"
    elif is_degraded:
        status["status"] = "degraded"
    else:
        status["status"] = "healthy"

    return status


# Tool registry for JSON-RPC
_tool_registry: dict[str, Any] = {}


def _python_type_to_json_schema(python_type: Any) -> dict[str, Any]:
    """Convert Python type annotation to JSON Schema."""
    type_name = getattr(python_type, "__name__", str(python_type))

    if python_type is str or type_name == "str":
        return {"type": "string"}
    elif python_type is int or type_name == "int":
        return {"type": "integer"}
    elif python_type is float or type_name == "float":
        return {"type": "number"}
    elif python_type is bool or type_name == "bool":
        return {"type": "boolean"}
    elif type_name == "NoneType":
        return {"type": "null"}
    elif hasattr(python_type, "__origin__"):
        origin = python_type.__origin__
        args = getattr(python_type, "__args__", ())

        if origin is list:
            item_type = args[0] if args else Any
            return {"type": "array", "items": _python_type_to_json_schema(item_type)}
        elif origin is dict:
            return {"type": "object"}
        elif origin is type(None) or str(origin) == "typing.Union":
            # Handle Optional types (Union with None)
            non_none_args = [a for a in args if a is not type(None)]
            if len(non_none_args) == 1:
                schema = _python_type_to_json_schema(non_none_args[0])
                schema["nullable"] = True
                return schema
            return {"type": "object"}
    return {"type": "object"}


def _get_tool_schema(func: Any) -> dict[str, Any]:
    """Extract JSON Schema from a tool function."""
    sig = inspect.signature(func)
    hints = get_type_hints(func) if hasattr(func, "__annotations__") else {}

    properties: dict[str, Any] = {}
    required: list[str] = []

    for name, param in sig.parameters.items():
        if name in ("self", "cls"):
            continue

        prop: dict[str, Any] = {}

        # Get type from hints
        if name in hints:
            prop = _python_type_to_json_schema(hints[name])

        # Get description and constraints from Field default (FieldInfo object)
        default_val = param.default
        if hasattr(default_val, "description") and default_val.description:
            prop["description"] = default_val.description
        if hasattr(default_val, "ge") and default_val.ge is not None:
            prop["minimum"] = default_val.ge
        if hasattr(default_val, "le") and default_val.le is not None:
            prop["maximum"] = default_val.le
        # Handle Field default value (check for PydanticUndefined)
        if hasattr(default_val, "default"):
            field_default = default_val.default
            # Check if it's the "required" sentinel (...)
            if field_default is not ... and not (
                hasattr(field_default, "__class__")
                and "PydanticUndefined" in field_default.__class__.__name__
            ):
                prop["default"] = field_default

        # Determine if required
        if param.default is inspect.Parameter.empty:
            required.append(name)
        elif hasattr(default_val, "default"):
            field_default = default_val.default
            # Required if default is ellipsis or PydanticUndefined
            if field_default is ... or (
                hasattr(field_default, "__class__")
                and "PydanticUndefined" in field_default.__class__.__name__
            ):
                required.append(name)

        properties[name] = prop

    return {
        "type": "object",
        "properties": properties,
        "required": required,
    }


def _register_tool(name: str, tool_or_func: Any, description: str | None = None) -> None:
    """Register a tool in the registry.

    Handles both raw functions and FastMCP FunctionTool objects.
    """
    # Extract the underlying function from FastMCP FunctionTool if needed
    if hasattr(tool_or_func, "fn"):
        func = tool_or_func.fn
        # Use FunctionTool's description if available
        if not description and hasattr(tool_or_func, "description") and tool_or_func.description:
            description = tool_or_func.description
    else:
        func = tool_or_func

    _tool_registry[name] = {
        "func": func,
        "description": description or (func.__doc__ or "").strip(),
        "schema": _get_tool_schema(func),
    }


@app.get("/mcp/tools")
async def list_mcp_tools() -> dict[str, Any]:
    """
    Return list of available MCP tools with their schemas.

    This endpoint enables tool discovery for the backend MCP client.
    """
    tools = []
    for name, info in _tool_registry.items():
        tools.append({
            "name": name,
            "description": info["description"],
            "inputSchema": info["schema"],
        })

    return {"tools": tools}


@app.post("/mcp")
async def mcp_rpc(request: Request) -> JSONResponse:
    """
    JSON-RPC 2.0 endpoint for MCP tool execution.

    Request format:
    {
        "jsonrpc": "2.0",
        "method": "<tool_name>",
        "params": {...},
        "id": <request_id>
    }

    Response format:
    {
        "jsonrpc": "2.0",
        "result": {...},
        "id": <request_id>
    }
    """
    try:
        body = await request.json()
    except Exception as e:
        return JSONResponse(
            status_code=400,
            content={
                "jsonrpc": "2.0",
                "error": {"code": -32700, "message": f"Parse error: {e}"},
                "id": None,
            },
        )

    # Validate JSON-RPC structure
    jsonrpc = body.get("jsonrpc")
    method = body.get("method")
    params = body.get("params", {})
    request_id = body.get("id")

    if jsonrpc != "2.0":
        return JSONResponse(
            status_code=400,
            content={
                "jsonrpc": "2.0",
                "error": {"code": -32600, "message": "Invalid Request: jsonrpc must be '2.0'"},
                "id": request_id,
            },
        )

    if not method:
        return JSONResponse(
            status_code=400,
            content={
                "jsonrpc": "2.0",
                "error": {"code": -32600, "message": "Invalid Request: method is required"},
                "id": request_id,
            },
        )

    # Look up tool
    tool_info = _tool_registry.get(method)
    if not tool_info:
        return JSONResponse(
            status_code=404,
            content={
                "jsonrpc": "2.0",
                "error": {"code": -32601, "message": f"Method not found: {method}"},
                "id": request_id,
            },
        )

    # Execute tool
    try:
        func = tool_info["func"]

        # Resolve Field defaults for missing parameters
        sig = inspect.signature(func)
        resolved_params = dict(params)
        for name, param in sig.parameters.items():
            if name not in resolved_params:
                default_val = param.default
                # Check if it's a FieldInfo with a default value
                if hasattr(default_val, "default"):
                    field_default = default_val.default
                    # Only use if it has an actual default (not required)
                    if field_default is not ... and not (
                        hasattr(field_default, "__class__")
                        and "PydanticUndefined" in field_default.__class__.__name__
                    ):
                        resolved_params[name] = field_default

        result = await func(**resolved_params)
        return JSONResponse(
            content={
                "jsonrpc": "2.0",
                "result": result,
                "id": request_id,
            }
        )
    except TypeError as e:
        return JSONResponse(
            status_code=400,
            content={
                "jsonrpc": "2.0",
                "error": {"code": -32602, "message": f"Invalid params: {e}"},
                "id": request_id,
            },
        )
    except Exception as e:
        logger.error(f"Tool execution error: {e}")
        return JSONResponse(
            status_code=500,
            content={
                "jsonrpc": "2.0",
                "error": {"code": -32000, "message": f"Server error: {e}"},
                "id": request_id,
            },
        )


# MCP Tools


@mcp.tool()
async def search_knowledge(
    project_id: str = Field(..., description="Project ID for scoping"),
    agent_id: str = Field(..., description="Agent ID making the request"),
    query: str = Field(..., description="Search query"),
    search_type: str = Field(
        default="hybrid",
        description="Search type: semantic, keyword, or hybrid",
    ),
    collection: str | None = Field(
        default=None,
        description="Collection to search (None = all)",
    ),
    limit: int = Field(
        default=10,
        ge=1,
        le=100,
        description="Maximum number of results",
    ),
    threshold: float = Field(
        default=0.7,
        ge=0.0,
        le=1.0,
        description="Minimum similarity score",
    ),
    file_types: list[str] | None = Field(
        default=None,
        description="Filter by file types (python, javascript, etc.)",
    ),
) -> dict[str, Any]:
    """
    Search the knowledge base for relevant content.

    Supports semantic (vector), keyword (full-text), and hybrid search.
    Returns chunks ranked by relevance to the query.
    """
    try:
        # Validate inputs
        if error := _validate_id(project_id, "project_id"):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if error := _validate_id(agent_id, "agent_id"):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if collection and (error := _validate_collection(collection)):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}

        # Parse search type
        try:
            search_type_enum = SearchType(search_type.lower())
        except ValueError:
            valid_types = [t.value for t in SearchType]
            return {
                "success": False,
                "error": f"Invalid search type: {search_type}. Valid types: {valid_types}",
                "code": ErrorCode.INVALID_REQUEST.value,
            }

        # Parse file types
        file_type_enums = None
        if file_types:
            try:
                file_type_enums = [FileType(ft.lower()) for ft in file_types]
            except ValueError as e:
                return {
                    "success": False,
                    "error": f"Invalid file type: {e}",
                    "code": ErrorCode.INVALID_REQUEST.value,
                }

        request = SearchRequest(
            project_id=project_id,
            agent_id=agent_id,
            query=query,
            search_type=search_type_enum,
            collection=collection,
            limit=limit,
            threshold=threshold,
            file_types=file_type_enums,
        )

        response = await _search.search(request)  # type: ignore[union-attr]

        return {
            "success": True,
            "query": response.query,
            "search_type": response.search_type,
            "results": [
                {
                    "id": r.id,
                    "content": r.content,
                    "score": r.score,
                    "source_path": r.source_path,
                    "start_line": r.start_line,
                    "end_line": r.end_line,
                    "chunk_type": r.chunk_type,
                    "file_type": r.file_type,
                    "collection": r.collection,
                    "metadata": r.metadata,
                }
                for r in response.results
            ],
            "total_results": response.total_results,
            "search_time_ms": response.search_time_ms,
        }

    except KnowledgeBaseError as e:
        logger.error(f"Search error: {e}")
        return {
            "success": False,
            "error": e.message,
            "code": e.code.value,
        }
    except Exception as e:
        logger.error(f"Unexpected search error: {e}")
        return {
            "success": False,
            "error": str(e),
            "code": ErrorCode.INTERNAL_ERROR.value,
        }


@mcp.tool()
async def ingest_content(
    project_id: str = Field(..., description="Project ID for scoping"),
    agent_id: str = Field(..., description="Agent ID making the request"),
    content: str = Field(..., description="Content to ingest"),
    source_path: str | None = Field(
        default=None,
        description="Source file path for reference",
    ),
    collection: str = Field(
        default="default",
        description="Collection to store in",
    ),
    chunk_type: str = Field(
        default="text",
        description="Content type: code, markdown, or text",
    ),
    file_type: str | None = Field(
        default=None,
        description="File type for code chunking (python, javascript, etc.)",
    ),
    metadata: dict[str, Any] | None = Field(
        default=None,
        description="Additional metadata to store",
    ),
) -> dict[str, Any]:
    """
    Ingest content into the knowledge base.

    Content is automatically chunked based on type, embedded using
    the LLM Gateway, and stored in pgvector for search.
    """
    try:
        # Validate inputs
        if error := _validate_id(project_id, "project_id"):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if error := _validate_id(agent_id, "agent_id"):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if error := _validate_collection(collection):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if error := _validate_source_path(source_path):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}

        # Validate content size to prevent DoS
        settings = get_settings()
        content_size = len(content.encode("utf-8"))
        if content_size > settings.max_document_size:
            return {
                "success": False,
                "error": f"Content size ({content_size} bytes) exceeds maximum allowed ({settings.max_document_size} bytes)",
                "code": ErrorCode.INVALID_REQUEST.value,
            }

        # Parse chunk type
        try:
            chunk_type_enum = ChunkType(chunk_type.lower())
        except ValueError:
            valid_types = [t.value for t in ChunkType]
            return {
                "success": False,
                "error": f"Invalid chunk type: {chunk_type}. Valid types: {valid_types}",
                "code": ErrorCode.INVALID_REQUEST.value,
            }

        # Parse file type
        file_type_enum = None
        if file_type:
            try:
                file_type_enum = FileType(file_type.lower())
            except ValueError:
                valid_types = [t.value for t in FileType]
                return {
                    "success": False,
                    "error": f"Invalid file type: {file_type}. Valid types: {valid_types}",
                    "code": ErrorCode.INVALID_REQUEST.value,
                }

        request = IngestRequest(
            project_id=project_id,
            agent_id=agent_id,
            content=content,
            source_path=source_path,
            collection=collection,
            chunk_type=chunk_type_enum,
            file_type=file_type_enum,
            metadata=metadata or {},
        )

        result = await _collections.ingest(request)  # type: ignore[union-attr]

        return {
            "success": result.success,
            "chunks_created": result.chunks_created,
            "embeddings_generated": result.embeddings_generated,
            "source_path": result.source_path,
            "collection": result.collection,
            "chunk_ids": result.chunk_ids,
            "error": result.error,
        }

    except KnowledgeBaseError as e:
        logger.error(f"Ingest error: {e}")
        return {
            "success": False,
            "error": e.message,
            "code": e.code.value,
        }
    except Exception as e:
        logger.error(f"Unexpected ingest error: {e}")
        return {
            "success": False,
            "error": str(e),
            "code": ErrorCode.INTERNAL_ERROR.value,
        }


@mcp.tool()
async def delete_content(
    project_id: str = Field(..., description="Project ID for scoping"),
    agent_id: str = Field(..., description="Agent ID making the request"),
    source_path: str | None = Field(
        default=None,
        description="Delete by source file path",
    ),
    collection: str | None = Field(
        default=None,
        description="Delete entire collection",
    ),
    chunk_ids: list[str] | None = Field(
        default=None,
        description="Delete specific chunk IDs",
    ),
) -> dict[str, Any]:
    """
    Delete content from the knowledge base.

    Specify either source_path, collection, or chunk_ids to delete.
    """
    try:
        # Validate inputs
        if error := _validate_id(project_id, "project_id"):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if error := _validate_id(agent_id, "agent_id"):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if collection and (error := _validate_collection(collection)):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if error := _validate_source_path(source_path):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}

        request = DeleteRequest(
            project_id=project_id,
            agent_id=agent_id,
            source_path=source_path,
            collection=collection,
            chunk_ids=chunk_ids,
        )

        result = await _collections.delete(request)  # type: ignore[union-attr]

        return {
            "success": result.success,
            "chunks_deleted": result.chunks_deleted,
            "error": result.error,
        }

    except KnowledgeBaseError as e:
        logger.error(f"Delete error: {e}")
        return {
            "success": False,
            "error": e.message,
            "code": e.code.value,
        }
    except Exception as e:
        logger.error(f"Unexpected delete error: {e}")
        return {
            "success": False,
            "error": str(e),
            "code": ErrorCode.INTERNAL_ERROR.value,
        }


@mcp.tool()
async def list_collections(
    project_id: str = Field(..., description="Project ID for scoping"),
    agent_id: str = Field(..., description="Agent ID making the request"),  # noqa: ARG001
) -> dict[str, Any]:
    """
    List all collections in a project's knowledge base.

    Returns collection names with chunk counts and file types.
    """
    try:
        # Validate inputs
        if error := _validate_id(project_id, "project_id"):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if error := _validate_id(agent_id, "agent_id"):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}

        result = await _collections.list_collections(project_id)  # type: ignore[union-attr]

        return {
            "success": True,
            "project_id": result.project_id,
            "collections": [
                {
                    "name": c.name,
                    "chunk_count": c.chunk_count,
                    "total_tokens": c.total_tokens,
                    "file_types": c.file_types,
                    "created_at": c.created_at.isoformat(),
                    "updated_at": c.updated_at.isoformat(),
                }
                for c in result.collections
            ],
            "total_collections": result.total_collections,
        }

    except KnowledgeBaseError as e:
        logger.error(f"List collections error: {e}")
        return {
            "success": False,
            "error": e.message,
            "code": e.code.value,
        }
    except Exception as e:
        logger.error(f"Unexpected list collections error: {e}")
        return {
            "success": False,
            "error": str(e),
            "code": ErrorCode.INTERNAL_ERROR.value,
        }


@mcp.tool()
async def get_collection_stats(
    project_id: str = Field(..., description="Project ID for scoping"),
    agent_id: str = Field(..., description="Agent ID making the request"),  # noqa: ARG001
    collection: str = Field(..., description="Collection name"),
) -> dict[str, Any]:
    """
    Get detailed statistics for a collection.

    Returns chunk counts, token totals, and type breakdowns.
    """
    try:
        # Validate inputs
        if error := _validate_id(project_id, "project_id"):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if error := _validate_id(agent_id, "agent_id"):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if error := _validate_collection(collection):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}

        stats = await _collections.get_collection_stats(project_id, collection)  # type: ignore[union-attr]

        return {
            "success": True,
            "collection": stats.collection,
            "project_id": stats.project_id,
            "chunk_count": stats.chunk_count,
            "unique_sources": stats.unique_sources,
            "total_tokens": stats.total_tokens,
            "avg_chunk_size": stats.avg_chunk_size,
            "chunk_types": stats.chunk_types,
            "file_types": stats.file_types,
            "oldest_chunk": stats.oldest_chunk.isoformat() if stats.oldest_chunk else None,
            "newest_chunk": stats.newest_chunk.isoformat() if stats.newest_chunk else None,
        }

    except KnowledgeBaseError as e:
        logger.error(f"Get collection stats error: {e}")
        return {
            "success": False,
            "error": e.message,
            "code": e.code.value,
        }
    except Exception as e:
        logger.error(f"Unexpected get collection stats error: {e}")
        return {
            "success": False,
            "error": str(e),
            "code": ErrorCode.INTERNAL_ERROR.value,
        }


@mcp.tool()
async def update_document(
    project_id: str = Field(..., description="Project ID for scoping"),
    agent_id: str = Field(..., description="Agent ID making the request"),
    source_path: str = Field(..., description="Source file path"),
    content: str = Field(..., description="New content"),
    collection: str = Field(
        default="default",
        description="Collection name",
    ),
    chunk_type: str = Field(
        default="text",
        description="Content type: code, markdown, or text",
    ),
    file_type: str | None = Field(
        default=None,
        description="File type for code chunking",
    ),
    metadata: dict[str, Any] | None = Field(
        default=None,
        description="Additional metadata",
    ),
) -> dict[str, Any]:
    """
    Update a document in the knowledge base.

    Replaces all existing chunks for the source path with new content.
    """
    try:
        # Validate inputs
        if error := _validate_id(project_id, "project_id"):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if error := _validate_id(agent_id, "agent_id"):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if error := _validate_collection(collection):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}
        if error := _validate_source_path(source_path):
            return {"success": False, "error": error, "code": ErrorCode.INVALID_REQUEST.value}

        # Validate content size to prevent DoS
        settings = get_settings()
        content_size = len(content.encode("utf-8"))
        if content_size > settings.max_document_size:
            return {
                "success": False,
                "error": f"Content size ({content_size} bytes) exceeds maximum allowed ({settings.max_document_size} bytes)",
                "code": ErrorCode.INVALID_REQUEST.value,
            }

        # Parse chunk type
        try:
            chunk_type_enum = ChunkType(chunk_type.lower())
        except ValueError:
            valid_types = [t.value for t in ChunkType]
            return {
                "success": False,
                "error": f"Invalid chunk type: {chunk_type}. Valid types: {valid_types}",
                "code": ErrorCode.INVALID_REQUEST.value,
            }

        # Parse file type
        file_type_enum = None
        if file_type:
            try:
                file_type_enum = FileType(file_type.lower())
            except ValueError:
                valid_types = [t.value for t in FileType]
                return {
                    "success": False,
                    "error": f"Invalid file type: {file_type}. Valid types: {valid_types}",
                    "code": ErrorCode.INVALID_REQUEST.value,
                }

        result = await _collections.update_document(  # type: ignore[union-attr]
            project_id=project_id,
            agent_id=agent_id,
            source_path=source_path,
            content=content,
            collection=collection,
            chunk_type=chunk_type_enum,
            file_type=file_type_enum,
            metadata=metadata,
        )

        return {
            "success": result.success,
            "chunks_created": result.chunks_created,
            "embeddings_generated": result.embeddings_generated,
            "source_path": result.source_path,
            "collection": result.collection,
            "chunk_ids": result.chunk_ids,
            "error": result.error,
        }

    except KnowledgeBaseError as e:
        logger.error(f"Update document error: {e}")
        return {
            "success": False,
            "error": e.message,
            "code": e.code.value,
        }
    except Exception as e:
        logger.error(f"Unexpected update document error: {e}")
        return {
            "success": False,
            "error": str(e),
            "code": ErrorCode.INTERNAL_ERROR.value,
        }


# Register tools in the JSON-RPC registry
# This must happen after tool functions are defined
_register_tool("search_knowledge", search_knowledge)
_register_tool("ingest_content", ingest_content)
_register_tool("delete_content", delete_content)
_register_tool("list_collections", list_collections)
_register_tool("get_collection_stats", get_collection_stats)
_register_tool("update_document", update_document)


def main() -> None:
    """Run the server."""
    import uvicorn

    settings = get_settings()

    uvicorn.run(
        "server:app",
        host=settings.host,
        port=settings.port,
        reload=settings.debug,
        log_level="info",
    )


if __name__ == "__main__":
    main()