108 lines
4.6 KiB
Python
108 lines
4.6 KiB
Python
"""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 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."""
|
|
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."""
|
|
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. Present accepted answers unchanged, retaining citations '
|
|
'and completeness notices. Never turn an error into a factual answer.'), **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])
|
|
server = create_server(stateless_http=True, json_response=True, transport_security=security)
|
|
|
|
@server.custom_route('/health', methods=['GET'])
|
|
async def bridge_health(request):
|
|
return JSONResponse({'status': 'ok', 'service': 'zpwiki-mcp'})
|
|
|
|
return BearerGate(server.streamable_http_app(), token)
|
|
|
|
|
|
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()
|