mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-30 20:35:22 +08:00
268 lines
11 KiB
Python
268 lines
11 KiB
Python
"""Component-owned, bounded async MCP listener; never uses the SDK singleton.
|
|
|
|
Fixed-port collisions fail rather than borrowing another invocation's listener.
|
|
The parent invocation revokes tokens and drains all accepted connections.
|
|
"""
|
|
|
|
import asyncio
|
|
import contextvars
|
|
import hmac
|
|
import json
|
|
import math
|
|
import secrets
|
|
import time
|
|
from contextvars import copy_context
|
|
|
|
from langbot_plugin.api.agent_tools.asset_gateway import (
|
|
LANGBOT_AGENT_GATEWAY_INFO,
|
|
LANGBOT_AGENT_GATEWAY_INSTRUCTIONS,
|
|
)
|
|
from langbot_plugin.api.agent_tools.external_tools import AgentRunExternalTools
|
|
from langbot_plugin.api.agent_tools.mcp_protocol import handle_mcp_payload, mcp_tool_error
|
|
from langbot_plugin.api.proxies.invocation import invocation_capability
|
|
|
|
MAX_BODY = 1024 * 1024
|
|
MAX_CONNECTIONS = 16
|
|
|
|
|
|
async def _create_tenant_free_listener(connected, host, port):
|
|
"""Create the shared listener from an explicit, empty Context.
|
|
|
|
``asyncio`` copies the current Context into the reader Handle that drives
|
|
``accept()``, and every accepted connection task inherits that Context. Built
|
|
inside the registering invocation, the handle would pin that invocation's
|
|
configuration and revocable authority to the shared listener for as long as
|
|
any other installation keeps using it. The listener is tenant-agnostic: an
|
|
invocation Context is entered only per validated run token, in
|
|
``Resolver.call_tool``.
|
|
"""
|
|
return await asyncio.get_running_loop().create_task(
|
|
asyncio.start_server(connected, host, port, limit=16384, backlog=16),
|
|
context=contextvars.Context(),
|
|
)
|
|
|
|
|
|
class Registration:
|
|
def __init__(self, gateway, api, ctx, ttl, authority, context):
|
|
self.context = context
|
|
self.authority = authority
|
|
self.gateway = gateway
|
|
self.token = secrets.token_urlsafe(32)
|
|
self.tools = AgentRunExternalTools(api, ctx)
|
|
self.expires_at = time.monotonic() + ttl
|
|
|
|
@property
|
|
def endpoint(self):
|
|
return self.gateway.endpoint
|
|
|
|
@property
|
|
def http_mcp_endpoint(self):
|
|
return self.endpoint + "/mcp"
|
|
|
|
async def stop(self):
|
|
await self.gateway.unregister(self.token)
|
|
|
|
|
|
class Gateway:
|
|
def __init__(self, timeout):
|
|
self.request_timeout = timeout
|
|
self.registrations = {}
|
|
self.server = None
|
|
self.tasks = set()
|
|
self.call_tasks = {}
|
|
self.pool = None
|
|
self.key = None
|
|
self.endpoint = ""
|
|
|
|
def _registration_for_token(self, token):
|
|
r = self.registrations.get(token)
|
|
if (
|
|
r is not None
|
|
and r.authority.active
|
|
and time.monotonic() < r.expires_at
|
|
and hmac.compare_digest(r.token, token)
|
|
):
|
|
return r
|
|
return None
|
|
|
|
def connected(self, reader, writer):
|
|
if len(self.tasks) >= MAX_CONNECTIONS:
|
|
writer.close()
|
|
return
|
|
# Serve connections tenant-free too; the registration's invocation
|
|
# context is entered only after a run token is validated.
|
|
task = asyncio.create_task(self.serve(reader, writer), context=contextvars.Context())
|
|
self.tasks.add(task)
|
|
task.add_done_callback(self.tasks.discard)
|
|
|
|
async def serve(self, reader, writer):
|
|
try:
|
|
async with asyncio.timeout(self.request_timeout):
|
|
head = await reader.readuntil(b"\r\n\r\n")
|
|
if len(head) > 16384:
|
|
raise ValueError("headers too large")
|
|
lines = head.decode("latin1").split("\r\n")
|
|
method, path, _ = lines[0].split(" ", 2)
|
|
headers = {}
|
|
for line in lines[1:]:
|
|
if not line:
|
|
continue
|
|
key, value = line.split(":", 1)
|
|
key = key.lower()
|
|
if key in headers:
|
|
raise ValueError("duplicate header")
|
|
headers[key] = value.strip()
|
|
if "transfer-encoding" in headers:
|
|
raise ValueError("chunked requests unsupported")
|
|
length = int(headers.get("content-length", "0"))
|
|
if not 0 <= length <= MAX_BODY:
|
|
await self.respond(writer, 413, {"error": "request too large"})
|
|
return
|
|
if method == "GET" and path == "/healthz":
|
|
await self.respond(writer, 200, {"ok": True})
|
|
return
|
|
if method != "POST" or path not in ("/mcp", "/mcp/http"):
|
|
await self.respond(writer, 404, {"error": "not found"})
|
|
return
|
|
payload = json.loads(await reader.readexactly(length))
|
|
if isinstance(payload, list) and len(payload) > 32:
|
|
raise ValueError("batch too large")
|
|
gateway = self
|
|
header_token = headers.get("authorization", "")
|
|
if header_token.lower().startswith("bearer "):
|
|
header_token = header_token[7:].strip()
|
|
else:
|
|
header_token = headers.get(
|
|
"x-langbot-agent-gateway-token", headers.get("x-langbot-agent-mcp-token", "")
|
|
)
|
|
|
|
class Resolver:
|
|
async def list_tools(self):
|
|
return AgentRunExternalTools.all_mcp_tools(include_run_token=True)
|
|
|
|
async def call_tool(self, name, arguments):
|
|
arguments = dict(arguments)
|
|
token = str(arguments.pop("run_token", "") or header_token)
|
|
registration = gateway._registration_for_token(token)
|
|
if registration is None:
|
|
return mcp_tool_error("A valid LangBot run_token is required")
|
|
request_task = asyncio.current_task()
|
|
gateway.call_tasks[request_task] = token
|
|
try:
|
|
registration.authority.require_active()
|
|
tool_task = registration.context.run(
|
|
asyncio.create_task,
|
|
registration.tools.call_mcp_tool(name, arguments),
|
|
)
|
|
try:
|
|
result = await tool_task
|
|
except asyncio.CancelledError:
|
|
tool_task.cancel()
|
|
await asyncio.gather(tool_task, return_exceptions=True)
|
|
raise
|
|
registration.authority.require_active()
|
|
return result
|
|
finally:
|
|
gateway.call_tasks.pop(request_task, None)
|
|
|
|
result = await handle_mcp_payload(
|
|
payload,
|
|
Resolver(),
|
|
server_info=LANGBOT_AGENT_GATEWAY_INFO,
|
|
instructions=LANGBOT_AGENT_GATEWAY_INSTRUCTIONS,
|
|
)
|
|
await self.respond(writer, 202 if result is None else 200, result)
|
|
except (ValueError, asyncio.LimitOverrunError, asyncio.IncompleteReadError):
|
|
pass
|
|
except (TimeoutError, ConnectionError):
|
|
pass
|
|
finally:
|
|
writer.close()
|
|
try:
|
|
await writer.wait_closed()
|
|
except (ConnectionError, OSError):
|
|
pass
|
|
|
|
async def respond(self, writer, status, result):
|
|
body = b"" if result is None else json.dumps(result, ensure_ascii=False).encode()
|
|
if len(body) > MAX_BODY:
|
|
status, body = 413, b'{"error":"response too large"}'
|
|
writer.write(
|
|
f"HTTP/1.1 {status} Response\r\nContent-Type: application/json\r\nContent-Length: {len(body)}\r\nConnection: close\r\n\r\n".encode()
|
|
+ body
|
|
)
|
|
await writer.drain()
|
|
|
|
async def unregister(self, token):
|
|
self.registrations.pop(token, None)
|
|
pending = [task for task, owner in self.call_tasks.items() if owner == token]
|
|
for task in pending:
|
|
task.cancel()
|
|
if not self.registrations:
|
|
# Remove from the pool before any await: no new lease on a closing listener.
|
|
if self.pool is not None and self.pool.get(self.key) is self:
|
|
del self.pool[self.key]
|
|
await self.stop()
|
|
elif pending:
|
|
await asyncio.gather(*pending, return_exceptions=True)
|
|
|
|
async def stop(self):
|
|
self.registrations.clear() # revoke before any await
|
|
if self.server is not None:
|
|
self.server.close()
|
|
for task in tuple(self.tasks):
|
|
task.cancel()
|
|
|
|
async def drain():
|
|
if self.server is not None:
|
|
await self.server.wait_closed()
|
|
await asyncio.gather(*tuple(self.tasks), return_exceptions=True)
|
|
|
|
cleanup = asyncio.create_task(drain())
|
|
cancelled = False
|
|
while not cleanup.done():
|
|
try:
|
|
await asyncio.shield(cleanup)
|
|
except asyncio.CancelledError:
|
|
cancelled = True
|
|
cleanup.result()
|
|
if cancelled:
|
|
raise asyncio.CancelledError
|
|
|
|
|
|
async def register_assets(owner, api, ctx, config):
|
|
timeout = float(config["asset_gateway_request_timeout"])
|
|
ttl = float(config["asset_gateway_token_ttl"])
|
|
if not math.isfinite(timeout) or not 0 < timeout <= 120:
|
|
raise ValueError("asset gateway request timeout must be in (0, 120]")
|
|
if not math.isfinite(ttl) or not 0 < ttl <= 3600:
|
|
raise ValueError("asset gateway token TTL must be in (0, 3600]")
|
|
# This owner is the SDK-created component, NOT a config-supplied tenant name.
|
|
# One component per artifact; registrations carry revocable invocation authority.
|
|
authority = invocation_capability(owner._plugin_runtime_handler)
|
|
if authority is None:
|
|
raise RuntimeError("Asset gateway requires an active run-scoped authority")
|
|
if not hasattr(owner, "_asset_gateway_pool"):
|
|
owner._asset_gateway_pool = {}
|
|
owner._asset_gateway_lock = asyncio.Lock()
|
|
pool = owner._asset_gateway_pool
|
|
key = (config["asset_gateway_host"], config["asset_gateway_port"], timeout)
|
|
async with owner._asset_gateway_lock:
|
|
gateway = pool.get(key)
|
|
if gateway is None:
|
|
gateway = Gateway(timeout)
|
|
gateway.pool, gateway.key = pool, key
|
|
try:
|
|
gateway.server = await _create_tenant_free_listener(gateway.connected, key[0], key[1])
|
|
host, port = gateway.server.sockets[0].getsockname()[:2]
|
|
host = f"[{host}]" if ":" in host else host
|
|
gateway.endpoint = f"http://{host}:{port}"
|
|
pool[key] = gateway
|
|
except BaseException:
|
|
await gateway.stop()
|
|
raise
|
|
authority.require_active()
|
|
registration = Registration(gateway, api, ctx, ttl, authority, copy_context())
|
|
gateway.registrations[registration.token] = registration
|
|
return registration
|