dp-zp-agent/packages/claude/mcp_server.py

"""MCP bridge: authenticated HTTP for Claude web, or local stdio for Desktop/Code."""
from __future__ import annotations

import argparse
import asyncio
import hmac
import json
import os
import sys
from pathlib import Path
from urllib.parse import urlsplit

sys.path.insert(0, str(Path(__file__).resolve().parent))
from client import ANSWER_PRESENTATION_POLICY, Config, backend_health, query_backend
from mcp.server.fastmcp import FastMCP
from mcp.server.transport_security import TransportSecuritySettings
from mcp.types import CallToolResult, TextContent, ToolAnnotations
from starlette.responses import JSONResponse


def tool_result(result):
    return CallToolResult(content=[TextContent(type='text', text=json.dumps(result, ensure_ascii=False))],
                          structuredContent=result, isError=result.get('status') == 'error')


async def zpwiki_rag(query: str, limit: int = 5) -> CallToolResult:
    """Verified ZP Wiki text answer. Keep the full question, names, work type and year.

    Follow presentation_instructions in the result; add no guesses or history commentary.
    """
    return tool_result(await asyncio.to_thread(query_backend, query, 'rag', limit=limit))


async def zpwiki_graphrag(query: str, limit: int = 5) -> CallToolResult:
    """Verified graph answer: relationships, shared topics, years and counts in ZP Wiki.

    Follow presentation_instructions in the result; no_answer means no supporting evidence.
    """
    return tool_result(await asyncio.to_thread(query_backend, query, 'graphrag', limit=limit))


async def zpwiki_health() -> CallToolResult:
    """Read readiness and data versions of both tools without calling a language model."""
    return tool_result(await asyncio.to_thread(backend_health))


def create_server(**kwargs):
    server = FastMCP('ZP Wiki', instructions=(
        'Use zpwiki_rag for text and individual facts; zpwiki_graphrag for relationships and counts. '
        'Tools return backend-verified answers. ' + ANSWER_PRESENTATION_POLICY), **kwargs)
    for handler in (zpwiki_rag, zpwiki_graphrag, zpwiki_health):
        server.add_tool(handler, annotations=ToolAnnotations(
            readOnlyHint=True, destructiveHint=False, openWorldHint=True))
    return server


class BearerGate:
    """Fixed connector credential, separate from backend/model keys. No OAuth simulation."""
    def __init__(self, app, token):
        self.app = app
        self.expected = ('Bearer ' + token).encode('ascii')

    async def __call__(self, scope, receive, send):
        if scope['type'] == 'http':
            headers = [value for key, value in scope['headers'] if key.lower() == b'authorization']
            if len(headers) != 1 or not hmac.compare_digest(headers[0], self.expected):
                response = JSONResponse({'error': 'unauthorized'}, status_code=401,
                                        headers={'WWW-Authenticate': 'Bearer'})
                await response(scope, receive, send)
                return
        await self.app(scope, receive, send)


def create_http_app():
    config = Config.load()
    token = os.getenv('ZP_MCP_ACCESS_TOKEN', '').strip()
    public_url = os.getenv('ZP_MCP_PUBLIC_URL', '').strip()
    url = urlsplit(public_url)
    url.port
    if (len(token) < 32 or not token.isascii() or any(c.isspace() for c in token)
            or hmac.compare_digest(token, config.key)
            or url.scheme != 'https' or not url.hostname or url.path != '/mcp'
            or url.username is not None or url.password is not None or url.query or url.fragment
            or any(c.isspace() for c in public_url)):
        raise ValueError('Configure a separate MCP token and an HTTPS public URL ending in /mcp')
    security = TransportSecuritySettings(enable_dns_rebinding_protection=True,
        allowed_hosts=[url.netloc, '127.0.0.1:*', 'localhost:*'],
        allowed_origins=['https://' + url.netloc])
    mode = os.getenv('ZP_MCP_AUTH_MODE', 'bearer')
    if mode not in {'bearer', 'oauth'}:
        raise ValueError('ZP_MCP_AUTH_MODE must be bearer or oauth')
    provider = None
    options = {}
    if mode == 'oauth':
        from oauth_provider import OAuthProvider, SCOPES, SCOPE
        from mcp.server.auth.settings import AuthSettings, ClientRegistrationOptions, RevocationOptions
        password = os.getenv('ZP_MCP_LOGIN_PASSWORD', '')
        if hmac.compare_digest(password.encode(), config.key.encode()):
            raise ValueError('The OAuth login code must differ from the backend key')
        provider = OAuthProvider(public_url, token, password,
            Path(os.getenv('ZP_MCP_OAUTH_DB', '/var/lib/zp-oauth/oauth.sqlite')))
        options = {'auth_server_provider': provider, 'auth': AuthSettings(
            issuer_url=provider.origin, resource_server_url=public_url, validate_token_resource=True,
            required_scopes=[SCOPE],
            client_registration_options=ClientRegistrationOptions(enabled=True, valid_scopes=SCOPES, default_scopes=SCOPES),
            revocation_options=RevocationOptions(enabled=True))}
    server = create_server(stateless_http=True, json_response=True, transport_security=security, **options)

    @server.custom_route('/health', methods=['GET'])
    async def bridge_health(request):
        if provider:
            authorization = request.headers.get('authorization', '')
            scheme, _, credential = authorization.partition(' ')
            if scheme.lower() != 'bearer' or not await provider.load_access_token(credential):
                return JSONResponse({'error':'unauthorized'}, status_code=401)
        return JSONResponse({'status': 'ok', 'service': 'zpwiki-mcp'})

    if not provider:
        return BearerGate(server.streamable_http_app(), token)
    server.custom_route('/oauth/login', methods=['GET','POST'])(provider.login)
    app = server.streamable_http_app()
    # SDK 1.x supports public clients but omits "none" from generated metadata.
    # Advertise the actually supported methods; leaving CIMD unadvertised selects DCR.
    from mcp.server.auth.routes import build_metadata
    metadata = build_metadata(server.settings.auth.issuer_url, None,
        server.settings.auth.client_registration_options, server.settings.auth.revocation_options)
    metadata.token_endpoint_auth_methods_supported = ['none','client_secret_post','client_secret_basic']

    async def discovery(request):
        return JSONResponse(metadata.model_dump(mode='json', exclude_none=True),
                            headers={'Access-Control-Allow-Origin':'*'})
    for route in app.routes:
        if getattr(route, 'path', None) == '/.well-known/oauth-authorization-server':
            route.endpoint = discovery
            from starlette.routing import request_response
            route.app = request_response(discovery)
    from oauth_provider import OAuthBoundary
    return OAuthBoundary(app, provider)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--transport', choices=['stdio', 'http'], default='stdio')
    args = parser.parse_args()
    if args.transport == 'stdio':
        create_server().run(transport='stdio')
    else:
        import uvicorn
        uvicorn.run(create_http_app(), host=os.getenv('ZP_MCP_BIND', '127.0.0.1'),
                    port=int(os.getenv('ZP_MCP_PORT', '8001')), access_log=False,
                    forwarded_allow_ips=os.getenv('FORWARDED_ALLOW_IPS', '127.0.0.1'))


if __name__ == '__main__':
    main()