Initial curated ITWorx_MCP_Hub source from 368fba691ea8d23403b4d56b53b5ec670d7c5020
This commit is contained in:
Public Source Publisher committed
commit
cb5be898ea
663 files changed
+82674
No files matched your search
@@ -0,0 +1,381 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
import yaml # type: ignore[import-untyped]
|
||||
from fastapi import FastAPI, Request, Response
|
||||
from sqlalchemy import select
|
||||
|
||||
from packages.hub_policy import Governance, PolicyEngine, PolicyRule, SlidingWindowRateLimiter
|
||||
from packages.hub_runtime import read_secret_file
|
||||
from services.control_api.auth import OIDCProviderClient, OIDCSettings, PrincipalMapper
|
||||
from services.control_api.auth.errors import InvalidToken, UnknownMachinePrincipal
|
||||
from services.control_api.auth.machine import MachineClientBinding
|
||||
from services.control_api.contextforge import ContextForgeClient, ContextForgeSettings
|
||||
from services.control_api.persistence import (
|
||||
SqlAlchemyAliasSource,
|
||||
create_engine_from_url,
|
||||
session_factory,
|
||||
)
|
||||
from services.control_api.persistence.demo_client import DEMO_CLIENT_ID
|
||||
from services.control_api.persistence.models import ClientRow, WorkspaceRow
|
||||
from services.dataplane_router import (
|
||||
AliasRegistry,
|
||||
DataplaneProxy,
|
||||
InvocationRecorder,
|
||||
RouterSettings,
|
||||
build_router,
|
||||
)
|
||||
from services.dataplane_router.models import InvocationMetadata
|
||||
|
||||
from .audit import SqlAlchemyAuditLog
|
||||
from .common import (
|
||||
hub_database_url,
|
||||
identity_assertion,
|
||||
load_connector_identity_key,
|
||||
run_asgi,
|
||||
)
|
||||
from .demo_token import TOKEN_PREFIX, DemoTokenAuthenticator, load_demo_token_signing_key
|
||||
from .mcp_skills import StaticSkills
|
||||
from .oauth_metadata import OAuthMetadataSettings, bearer_challenge
|
||||
|
||||
|
||||
class JsonLineInvocationRecorder(InvocationRecorder):
|
||||
"""Emit bounded response-contract metadata without request data or credentials."""
|
||||
|
||||
def record(self, metadata: InvocationMetadata) -> None:
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"event": "mcp.response.completed",
|
||||
"traceId": metadata.trace_id,
|
||||
"project": metadata.project_slug,
|
||||
"mcpMethod": metadata.mcp_method,
|
||||
"mcpName": metadata.mcp_name,
|
||||
"protocolRevision": metadata.protocol_revision,
|
||||
"responseStatus": metadata.response_status,
|
||||
"responseBytes": metadata.response_bytes,
|
||||
"responseContentType": metadata.response_content_type,
|
||||
"jsonrpcResponse": metadata.jsonrpc_response,
|
||||
"responseIdMatches": metadata.response_id_matches,
|
||||
"responseShape": metadata.response_shape,
|
||||
"durationMs": metadata.duration_ms,
|
||||
},
|
||||
separators=(",", ":"),
|
||||
sort_keys=True,
|
||||
),
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
def _policy_engine(catalog_root: Path) -> PolicyEngine:
|
||||
policies: list[PolicyRule] = []
|
||||
for path in sorted((catalog_root / "policies").glob("*.yaml")):
|
||||
document = yaml.safe_load(path.read_text(encoding="utf-8"))
|
||||
if not isinstance(document, dict):
|
||||
raise ValueError("policy catalog document is invalid")
|
||||
policies.append(PolicyRule.from_document(document))
|
||||
if not policies or not any(policy.slug == "global-deny" for policy in policies):
|
||||
raise ValueError("policy catalog must include global-deny")
|
||||
return PolicyEngine(policies)
|
||||
|
||||
|
||||
def _tool_governance_metadata(catalog_root: Path) -> dict[str, tuple[str, str, str]]:
|
||||
metadata: dict[str, tuple[str, str, str]] = {}
|
||||
for path in sorted((catalog_root / "tools").glob("*.yaml")):
|
||||
document = yaml.safe_load(path.read_text(encoding="utf-8"))
|
||||
if not isinstance(document, dict):
|
||||
raise ValueError("tool catalog document is invalid")
|
||||
spec = document.get("spec")
|
||||
if not isinstance(spec, dict):
|
||||
raise ValueError("tool catalog spec is invalid")
|
||||
name = spec.get("canonicalName")
|
||||
connector = spec.get("connectorType")
|
||||
risk_class = spec.get("riskClass")
|
||||
if not isinstance(name, str) or not name:
|
||||
raise ValueError("tool governance metadata is incomplete")
|
||||
if not isinstance(connector, str) or not connector:
|
||||
raise ValueError("tool governance metadata is incomplete")
|
||||
if not isinstance(risk_class, str) or not risk_class:
|
||||
raise ValueError("tool governance metadata is incomplete")
|
||||
published_name = f"itworx-{connector}-{name.replace('.', '-')}"
|
||||
configured_names = spec.get("exposedNames", [])
|
||||
if not isinstance(configured_names, list) or not all(
|
||||
isinstance(item, str) for item in configured_names
|
||||
):
|
||||
raise ValueError("tool governance exposed names are invalid")
|
||||
for exposed_name in (name, published_name, *configured_names):
|
||||
if exposed_name in metadata:
|
||||
raise ValueError(f"duplicate tool governance metadata: {exposed_name}")
|
||||
metadata[exposed_name] = (connector, risk_class, name)
|
||||
if not metadata:
|
||||
raise ValueError("tool governance metadata is empty")
|
||||
return metadata
|
||||
|
||||
|
||||
class DatabaseAliasSource:
|
||||
def __init__(self, sessions: Any, gateway_url: str) -> None:
|
||||
self.sessions = sessions
|
||||
self.gateway_url = gateway_url
|
||||
|
||||
def aliases(self) -> tuple[Any, ...]:
|
||||
with self.sessions() as session:
|
||||
workspace_id = session.scalar(
|
||||
select(WorkspaceRow.id).where(WorkspaceRow.slug == "itworx")
|
||||
)
|
||||
if workspace_id is None:
|
||||
return ()
|
||||
return SqlAlchemyAliasSource(
|
||||
session,
|
||||
workspace_id=workspace_id,
|
||||
gateway_base_url=self.gateway_url,
|
||||
).aliases()
|
||||
|
||||
|
||||
class DatabaseMachineAuthenticator:
|
||||
_bearer = re.compile(r"^Bearer ([A-Za-z0-9._~+/=-]{1,65536})$")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
providers: OIDCProviderClient | tuple[OIDCProviderClient, ...],
|
||||
sessions: Any,
|
||||
) -> None:
|
||||
if not isinstance(providers, tuple):
|
||||
providers = (providers,)
|
||||
if not providers:
|
||||
raise ValueError("at least one OIDC provider is required")
|
||||
self.providers = providers
|
||||
self.sessions = sessions
|
||||
|
||||
def __call__(self, authorization_header: str | None) -> Any:
|
||||
if authorization_header is None:
|
||||
raise InvalidToken()
|
||||
match = self._bearer.fullmatch(authorization_header.strip())
|
||||
if match is None:
|
||||
raise InvalidToken()
|
||||
claims = None
|
||||
for provider in self.providers:
|
||||
try:
|
||||
claims = provider.validate_token(match.group(1))
|
||||
break
|
||||
except InvalidToken:
|
||||
continue
|
||||
if claims is None:
|
||||
raise InvalidToken()
|
||||
with self.sessions() as session:
|
||||
client = session.scalar(
|
||||
select(ClientRow).where(
|
||||
ClientRow.issuer == claims.issuer,
|
||||
ClientRow.subject == claims.subject,
|
||||
ClientRow.status == "active",
|
||||
)
|
||||
)
|
||||
if client is None:
|
||||
raise UnknownMachinePrincipal()
|
||||
client.last_seen_at = datetime.now(UTC)
|
||||
session.commit()
|
||||
binding = MachineClientBinding(
|
||||
subject=client.subject,
|
||||
client_id=UUID(str(client.id)),
|
||||
enabled=True,
|
||||
)
|
||||
return PrincipalMapper(machine_subjects={binding.subject: binding}).map_claims(claims)
|
||||
|
||||
|
||||
class CompositeAuthenticator:
|
||||
"""Routes demo-token-shaped bearer credentials to the demo path, else OIDC."""
|
||||
|
||||
def __init__(
|
||||
self, demo: Callable[[str | None], Any], oidc: Callable[[str | None], Any]
|
||||
) -> None:
|
||||
self.demo = demo
|
||||
self.oidc = oidc
|
||||
|
||||
def __call__(self, authorization_header: str | None) -> Any:
|
||||
if authorization_header is not None and authorization_header.startswith(
|
||||
f"Bearer {TOKEN_PREFIX}"
|
||||
):
|
||||
return self.demo(authorization_header)
|
||||
return self.oidc(authorization_header)
|
||||
|
||||
|
||||
class DemoAwareRateLimiter(SlidingWindowRateLimiter):
|
||||
"""Applies a stricter shared ceiling to the one fixed demo client identity.
|
||||
|
||||
Every other client_id behaves exactly as the base SlidingWindowRateLimiter
|
||||
it subclasses — this only special-cases the demo client so Governance's
|
||||
`rate_limiter: SlidingWindowRateLimiter | None` type still holds.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
default_limit: int,
|
||||
demo_limit: int,
|
||||
demo_client_id: UUID,
|
||||
window_seconds: float = 60.0,
|
||||
) -> None:
|
||||
super().__init__(limit=default_limit, window_seconds=window_seconds)
|
||||
self._demo_limiter = SlidingWindowRateLimiter(
|
||||
limit=demo_limit, window_seconds=window_seconds
|
||||
)
|
||||
self._demo_client_id = demo_client_id
|
||||
|
||||
def allow(self, *, client_id: UUID | str, project: str, tool: str) -> bool:
|
||||
if str(client_id) == str(self._demo_client_id):
|
||||
return self._demo_limiter.allow(client_id=client_id, project=project, tool=tool)
|
||||
return super().allow(client_id=client_id, project=project, tool=tool)
|
||||
|
||||
|
||||
def _oidc_providers() -> tuple[OIDCProviderClient, ...]:
|
||||
providers = [OIDCProviderClient(OIDCSettings.from_env())]
|
||||
additional = {
|
||||
"issuer": os.environ.get("HUB_MCP_ADDITIONAL_OIDC_ISSUER", "").strip(),
|
||||
"audience": os.environ.get("HUB_MCP_ADDITIONAL_OIDC_AUDIENCE", "").strip(),
|
||||
"client_id": os.environ.get("HUB_MCP_ADDITIONAL_OIDC_CLIENT_ID", "").strip(),
|
||||
"redirect_uri": os.environ.get("HUB_MCP_ADDITIONAL_OIDC_REDIRECT_URI", "").strip(),
|
||||
}
|
||||
configured = [bool(value) for value in additional.values()]
|
||||
if any(configured) and not all(configured):
|
||||
raise ValueError("additional MCP OIDC provider configuration is incomplete")
|
||||
if all(configured):
|
||||
settings = OIDCSettings.from_env(
|
||||
{
|
||||
"HUB_OIDC_ISSUER": additional["issuer"],
|
||||
"HUB_OIDC_AUDIENCE": additional["audience"],
|
||||
"HUB_OIDC_CLIENT_ID": additional["client_id"],
|
||||
"HUB_OIDC_REDIRECT_URI": additional["redirect_uri"],
|
||||
"HUB_OIDC_ALLOWED_ALGORITHMS": os.environ.get(
|
||||
"HUB_OIDC_ALLOWED_ALGORITHMS", "RS256"
|
||||
),
|
||||
}
|
||||
)
|
||||
providers.append(OIDCProviderClient(settings))
|
||||
return tuple(providers)
|
||||
|
||||
|
||||
def build_app() -> FastAPI:
|
||||
oauth_metadata = OAuthMetadataSettings.from_environment()
|
||||
engine = create_engine_from_url(hub_database_url())
|
||||
sessions = session_factory(engine)
|
||||
oidc_providers = _oidc_providers()
|
||||
oidc_provider = oidc_providers[0]
|
||||
contextforge = ContextForgeClient(ContextForgeSettings.from_env())
|
||||
catalog_root = Path(os.environ.get("HUB_CATALOG_ROOT", "/opt/hub/catalog")).resolve()
|
||||
rate_limit = int(os.environ.get("HUB_RATE_LIMIT_PER_MINUTE", "600"))
|
||||
if not 1 <= rate_limit <= 10_000:
|
||||
raise ValueError("HUB_RATE_LIMIT_PER_MINUTE must be between 1 and 10000")
|
||||
demo_rate_limit = int(os.environ.get("DEMO_TOKEN_RATE_LIMIT_PER_MINUTE", "10"))
|
||||
if not 1 <= demo_rate_limit <= rate_limit:
|
||||
raise ValueError(
|
||||
"DEMO_TOKEN_RATE_LIMIT_PER_MINUTE must be between 1 and HUB_RATE_LIMIT_PER_MINUTE"
|
||||
)
|
||||
governance = Governance(
|
||||
_policy_engine(catalog_root),
|
||||
SqlAlchemyAuditLog(sessions, issuer=str(oidc_provider.settings.issuer)),
|
||||
rate_limiter=DemoAwareRateLimiter(
|
||||
default_limit=rate_limit, demo_limit=demo_rate_limit, demo_client_id=DEMO_CLIENT_ID
|
||||
),
|
||||
)
|
||||
connector_identity_key = load_connector_identity_key()
|
||||
intake_key_path = os.environ.get("KNOWLEDGE_HUB_INGESTION_IDENTITY_KEY_FILE")
|
||||
intake_identity_key: bytes | None = None
|
||||
if intake_key_path:
|
||||
intake_identity_key = read_secret_file(Path(intake_key_path)).encode("utf-8")
|
||||
if len(intake_identity_key) < 32:
|
||||
raise ValueError("knowledge intake identity key must contain at least 32 bytes")
|
||||
if (
|
||||
os.environ.get("INGESTION_CONNECTOR_ENABLED", "false").lower() == "true"
|
||||
and not intake_identity_key
|
||||
):
|
||||
raise ValueError("knowledge intake identity key is required when intake is enabled")
|
||||
alias_cache_ttl = int(os.environ.get("HUB_ALIAS_CACHE_TTL_SECONDS", "5"))
|
||||
if not 1 <= alias_cache_ttl <= 60:
|
||||
raise ValueError("HUB_ALIAS_CACHE_TTL_SECONDS must be between 1 and 60")
|
||||
router_settings = RouterSettings(
|
||||
alias_cache_ttl_seconds=alias_cache_ttl,
|
||||
max_upstream_concurrency=int(os.environ.get("HUB_UPSTREAM_MAX_CONCURRENCY", "24")),
|
||||
upstream_queue_timeout_seconds=float(
|
||||
os.environ.get("HUB_UPSTREAM_QUEUE_TIMEOUT_SECONDS", "5")
|
||||
),
|
||||
)
|
||||
# Use ContextForge's unique control-network alias. The generic alias is
|
||||
# present on several Docker networks and can resolve outside the service
|
||||
# token's source-IP allowlist, causing intermittent 403 responses.
|
||||
source = DatabaseAliasSource(sessions, "http://contextforge-control:4444/mcp")
|
||||
registry = AliasRegistry(source, settings=router_settings)
|
||||
demo_revocations_path = Path(
|
||||
os.environ.get("DEMO_TOKEN_REVOCATIONS_FILE", "/etc/itworx/demo-token-revocations.json")
|
||||
)
|
||||
authenticator = CompositeAuthenticator(
|
||||
demo=DemoTokenAuthenticator(
|
||||
None, demo_revocations_path, key_loader=load_demo_token_signing_key
|
||||
),
|
||||
oidc=DatabaseMachineAuthenticator(oidc_providers, sessions),
|
||||
)
|
||||
proxy = DataplaneProxy(
|
||||
registry,
|
||||
authenticator,
|
||||
recorder=JsonLineInvocationRecorder(),
|
||||
settings=router_settings,
|
||||
governance=governance,
|
||||
tool_metadata=_tool_governance_metadata(catalog_root),
|
||||
upstream_authorization=contextforge.authorization_header,
|
||||
identity_assertion=lambda project, environment, client_id, trace_id: identity_assertion(
|
||||
(intake_identity_key or b"") if project == "knowledge-hub" else connector_identity_key,
|
||||
project=project,
|
||||
environment=environment,
|
||||
client_id=client_id,
|
||||
trace_id=trace_id,
|
||||
),
|
||||
mcp_extension=StaticSkills(serve_public_widgets=True),
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: FastAPI) -> AsyncIterator[None]:
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await proxy.close()
|
||||
for provider in oidc_providers:
|
||||
provider.close()
|
||||
contextforge.close()
|
||||
engine.dispose()
|
||||
|
||||
app = FastAPI(
|
||||
title="ITWorx MCP Hub dataplane",
|
||||
docs_url=None,
|
||||
redoc_url=None,
|
||||
openapi_url=None,
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
@app.middleware("http")
|
||||
async def advertise_chatgpt_oauth(
|
||||
request: Request, call_next: Callable[[Request], Awaitable[Response]]
|
||||
) -> Response:
|
||||
response = await call_next(request)
|
||||
path_parts = request.url.path.removeprefix("/mcp/").split("/")
|
||||
if (
|
||||
len(path_parts) == 2
|
||||
and all(part.replace("-", "").isalnum() for part in path_parts)
|
||||
and response.status_code == 401
|
||||
):
|
||||
resource_path = "/".join(path_parts)
|
||||
response.headers["WWW-Authenticate"] = bearer_challenge(resource_path, oauth_metadata)
|
||||
return response
|
||||
|
||||
app.include_router(build_router(proxy))
|
||||
return app
|
||||
|
||||
|
||||
def run() -> None:
|
||||
run_asgi(build_app())
|
||||
Reference in new issue
Block a user