wehub-resource-sync
328 lines
14 KiB
Python
328 lines
14 KiB
Python
"""Runner-side MCP proxy manager.
|
|
|
|
Routes all MCP calls through the Omnigent server's
|
|
``POST /v1/sessions/{session_id}/mcp`` endpoint (MCP Streamable HTTP,
|
|
JSON-RPC 2.0) instead of connecting to external MCP servers directly.
|
|
|
|
The Omnigent server holds the live connections in its
|
|
:class:`omnigent.server.mcp_pool.ServerMcpPool` and enforces
|
|
TOOL_CALL + TOOL_RESULT policies on every call before forwarding to the
|
|
real MCP server. :class:`ProxyMcpManager` implements the same public
|
|
interface as :class:`omnigent.runner.mcp_manager.RunnerMcpManager` so
|
|
it can be substituted transparently at every dispatch site in
|
|
``runner/app.py``.
|
|
|
|
When to use each implementation:
|
|
|
|
- **AP mode** (out-of-process runner paired with Omnigent server):
|
|
``ProxyMcpManager(session_id, server_client)`` — all tool calls flow
|
|
through the Omnigent server, policy enforced centrally.
|
|
- **No-AP / test mode**:
|
|
``RunnerMcpManager(stdio_cwd=...)`` — direct MCP connections, runner
|
|
enforces policy locally via :class:`RunnerToolPolicyGate`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from omnigent.runner import pending_approvals
|
|
from omnigent.runner.mcp_manager import McpSchemasResult
|
|
from omnigent.runner.tool_dispatch import MCP_PROXY_CALL_TIMEOUT_S
|
|
from omnigent.spec.types import AgentSpec
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ProxyMcpManager:
|
|
"""Routes runner-side MCP calls through the Omnigent server MCP proxy endpoint.
|
|
|
|
Drop-in substitute for :class:`omnigent.runner.mcp_manager.RunnerMcpManager`
|
|
used in Omnigent mode. The runner creates one instance per session on first MCP use
|
|
and passes it to :func:`omnigent.runner.tool_dispatch.execute_tool` in place
|
|
of a direct manager.
|
|
|
|
:param session_id: The AP-allocated session (conversation) id, e.g.
|
|
``"conv_abc123"``. Used to build the proxy endpoint URL:
|
|
``/v1/sessions/{session_id}/mcp``.
|
|
:param ap_client: An :class:`httpx.AsyncClient` pointed at the Omnigent server.
|
|
Must already carry the runner's service auth (e.g. Databricks bearer
|
|
token) so requests are accepted by the Omnigent server's auth middleware.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
session_id: str,
|
|
ap_client: httpx.AsyncClient,
|
|
publish_event: Callable[[str, dict[str, Any]], None] | None = None,
|
|
) -> None:
|
|
"""Create a proxy manager bound to one session.
|
|
|
|
:param session_id: Omnigent session id, e.g. ``"conv_abc123"``.
|
|
:param ap_client: HTTP client pointed at the Omnigent server.
|
|
:param publish_event: Optional callback that puts an SSE event on
|
|
the runner's per-session outbound queue. Required for the
|
|
approval flow so ``response.elicitation_resolved`` is emitted
|
|
when the user decides (keeps the approval-badge counter in
|
|
sync). Pass ``None`` only in test contexts where the badge
|
|
is irrelevant.
|
|
"""
|
|
self._session_id = session_id
|
|
self._omnigent_client = ap_client
|
|
self._publish_event = publish_event
|
|
|
|
@property
|
|
def _mcp_url(self) -> str:
|
|
"""MCP proxy endpoint URL for this session.
|
|
|
|
:returns: Path string, e.g. ``"/v1/sessions/conv_abc123/mcp"``.
|
|
"""
|
|
return f"/v1/sessions/{self._session_id}/mcp"
|
|
|
|
async def schemas_for(self, spec: AgentSpec) -> McpSchemasResult:
|
|
"""Fetch tool schemas from the Omnigent server MCP proxy (``tools/list``).
|
|
|
|
Sends a ``tools/list`` JSON-RPC 2.0 request to the Omnigent server's MCP
|
|
proxy endpoint. Returns tool schemas in the same flat OpenAI
|
|
function-tool format as
|
|
:meth:`omnigent.runner.mcp_manager.RunnerMcpManager.schemas_for`.
|
|
|
|
Tool names are returned with the server namespace prefix applied by
|
|
the Omnigent server (e.g. ``github__search``). This ensures the harness
|
|
sees collision-safe names even when multiple MCP servers define tools
|
|
with the same bare name.
|
|
|
|
:param spec: The agent spec. When ``spec.mcp_servers`` is empty,
|
|
returns an empty result immediately without hitting the network.
|
|
:returns: :class:`McpSchemasResult` containing schemas, tool name
|
|
set, and per-server failure messages.
|
|
"""
|
|
if not spec.mcp_servers:
|
|
return McpSchemasResult(schemas=[], tool_names=set(), failures={})
|
|
|
|
payload: dict[str, Any] = {
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "tools/list",
|
|
"params": {},
|
|
}
|
|
try:
|
|
resp = await self._omnigent_client.post(
|
|
self._mcp_url,
|
|
json=payload,
|
|
timeout=30.0,
|
|
)
|
|
resp.raise_for_status()
|
|
data: dict[str, Any] = resp.json()
|
|
except Exception as exc: # noqa: BLE001 — network + HTTP errors all surface as failures
|
|
_logger.warning(
|
|
"ProxyMcpManager tools/list failed for session %r: %s",
|
|
self._session_id,
|
|
exc,
|
|
)
|
|
return McpSchemasResult(
|
|
schemas=[],
|
|
tool_names=set(),
|
|
failures={"proxy": f"{type(exc).__name__}: {exc}"},
|
|
)
|
|
|
|
if "error" in data:
|
|
err: dict[str, Any] = data["error"]
|
|
msg = f"MCP proxy error {err.get('code')}: {err.get('message')}"
|
|
_logger.warning(
|
|
"ProxyMcpManager tools/list returned RPC error for session %r: %s",
|
|
self._session_id,
|
|
msg,
|
|
)
|
|
return McpSchemasResult(
|
|
schemas=[],
|
|
tool_names=set(),
|
|
failures={"proxy": msg},
|
|
)
|
|
|
|
tools_list: list[dict[str, Any]] = data.get("result", {}).get("tools", [])
|
|
schemas: list[dict[str, Any]] = []
|
|
tool_names: set[str] = set()
|
|
for tool in tools_list:
|
|
name = tool.get("name", "")
|
|
if not name:
|
|
continue
|
|
# The Omnigent server returns ``inputSchema`` (JSON Schema from MCP).
|
|
# Convert to the ``parameters`` key expected by LLM providers,
|
|
# normalizing the same way RunnerMcpManager does via
|
|
# _normalize_input_schema (ensure ``properties`` key is present).
|
|
raw_schema: dict[str, Any] | None = tool.get("inputSchema")
|
|
parameters: dict[str, Any]
|
|
if raw_schema is None:
|
|
parameters = {"type": "object", "properties": {}}
|
|
elif raw_schema.get("type") == "object" and "properties" not in raw_schema:
|
|
parameters = {**raw_schema, "properties": {}}
|
|
else:
|
|
parameters = raw_schema
|
|
schema: dict[str, Any] = {
|
|
"type": "function",
|
|
"name": name,
|
|
"description": tool.get("description", ""),
|
|
"parameters": parameters,
|
|
}
|
|
schemas.append(schema)
|
|
tool_names.add(name)
|
|
|
|
return McpSchemasResult(schemas=schemas, tool_names=tool_names, failures={})
|
|
|
|
async def call_tool(
|
|
self,
|
|
spec: AgentSpec | None,
|
|
tool_name: str,
|
|
# Values are Any because MCP tool arguments are JSON objects with
|
|
# heterogeneous value types (str, int, bool, nested dicts, etc.).
|
|
arguments: dict[str, Any], # type: ignore[explicit-any]
|
|
) -> str:
|
|
"""Dispatch a tool call via the Omnigent server MCP proxy (``tools/call``).
|
|
|
|
Sends a ``tools/call`` JSON-RPC 2.0 request. The Omnigent server enforces
|
|
TOOL_CALL and TOOL_RESULT policies before forwarding to the real MCP
|
|
server.
|
|
|
|
**ASK policy / approval flow**: when the Omnigent server returns an
|
|
``InputRequiredResult`` (MCP MRTR spec), this method parks on the
|
|
runner-side approval Future until the user accepts or declines, then
|
|
retries once with the user's decision in ``inputResponses``.
|
|
|
|
:param spec: Ignored — accepted for interface parity with
|
|
:class:`RunnerMcpManager`. ``None`` is acceptable for callers
|
|
that have no spec context (e.g. the claude-native relay executor).
|
|
:param tool_name: Tool name as seen by the LLM, e.g.
|
|
``"github__search"`` or ``"sys_os_read"``.
|
|
:param arguments: Decoded tool arguments dict.
|
|
:returns: Tool output string. On denial the Omnigent server returns an
|
|
MCP error response which is converted here to a JSON error string
|
|
so the harness can feed it to the LLM as a tool result.
|
|
:raises RuntimeError: On network failure or unexpected protocol errors.
|
|
"""
|
|
del spec # Omnigent server resolves spec from session context
|
|
|
|
payload: dict[str, Any] = {
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "tools/call",
|
|
"params": {"name": tool_name, "arguments": arguments},
|
|
}
|
|
|
|
# At most two iterations: initial call + one approval retry.
|
|
for _attempt in range(2):
|
|
try:
|
|
resp = await self._omnigent_client.post(
|
|
self._mcp_url,
|
|
json=payload,
|
|
# Short connect timeout still fails fast on an unreachable
|
|
# server; the read timeout covers ordinary proxy request
|
|
# hangs. Sub-agent dispatch returns an async handle
|
|
# immediately and no longer holds this call for a child turn.
|
|
timeout=httpx.Timeout(
|
|
connect=10.0,
|
|
read=MCP_PROXY_CALL_TIMEOUT_S,
|
|
write=10.0,
|
|
pool=10.0,
|
|
),
|
|
)
|
|
resp.raise_for_status()
|
|
data: dict[str, Any] = resp.json()
|
|
except Exception as exc:
|
|
raise RuntimeError(
|
|
f"MCP proxy call failed for tool {tool_name!r} in session "
|
|
f"{self._session_id!r}: {exc}"
|
|
) from exc
|
|
|
|
if "error" in data:
|
|
err: dict[str, Any] = data["error"]
|
|
code = err.get("code")
|
|
msg = err.get("message", "")
|
|
# -32000 is the MCP convention for server-defined errors (tool
|
|
# denials, tool errors). Return as a JSON error string so the
|
|
# harness feeds the refusal back to the LLM rather than raising.
|
|
if code == -32000:
|
|
return json.dumps({"error": msg})
|
|
raise RuntimeError(
|
|
f"MCP proxy protocol error {code} for tool {tool_name!r}: {msg}"
|
|
)
|
|
|
|
result: dict[str, Any] = data.get("result", {})
|
|
|
|
# ── ASK: Omnigent server returned InputRequiredResult ───────────────
|
|
# Park for user approval and retry with inputResponses per the
|
|
# MCP Multi Round-Trip Requests spec.
|
|
if result.get("resultType") == "input_required":
|
|
if _attempt >= 1:
|
|
# Guard against unexpected re-elicitation after one retry.
|
|
return json.dumps({"error": "Approval loop exceeded"})
|
|
|
|
input_requests: dict[str, Any] = result.get("inputRequests") or {}
|
|
request_state: str = result.get("requestState", "")
|
|
# The Omnigent server uses the elicitation_id as the key in
|
|
# inputRequests (MRTR spec: keys are server-assigned and the
|
|
# client CAN read them — only requestState is opaque).
|
|
elicitation_id = next(iter(input_requests), "")
|
|
if not elicitation_id:
|
|
return json.dumps(
|
|
{"error": "Approval required but no elicitation in inputRequests"}
|
|
)
|
|
|
|
_pub: Callable[[str, dict[str, Any]], None] = (
|
|
self._publish_event
|
|
if self._publish_event is not None
|
|
else (lambda _s, _e: None)
|
|
)
|
|
approved = await pending_approvals.wait_for_user_approval(
|
|
elicitation_id=elicitation_id,
|
|
conversation_id=self._session_id,
|
|
publish_event=_pub,
|
|
)
|
|
|
|
payload = {
|
|
"jsonrpc": "2.0",
|
|
"id": 2, # MRTR spec: retry MUST use a different id
|
|
"method": "tools/call",
|
|
"params": {
|
|
"name": tool_name,
|
|
"arguments": arguments,
|
|
"requestState": request_state,
|
|
"inputResponses": {
|
|
elicitation_id: {"action": "accept" if approved else "decline"}
|
|
},
|
|
},
|
|
}
|
|
continue
|
|
|
|
# ── Normal result (ALLOW or post-approval execution) ──────────
|
|
content = result.get("content", [])
|
|
is_error = result.get("isError", False)
|
|
if isinstance(content, list):
|
|
parts = [
|
|
block.get("text", "")
|
|
for block in content
|
|
if isinstance(block, dict) and block.get("type") == "text"
|
|
]
|
|
text = "\n".join(p for p in parts if p)
|
|
if is_error:
|
|
return json.dumps({"error": text})
|
|
return text if text else json.dumps(result)
|
|
return json.dumps(result)
|
|
|
|
return json.dumps({"error": "Approval retry loop exhausted"})
|
|
|
|
async def prewarm(self, spec: AgentSpec) -> None:
|
|
"""No-op — the Omnigent server warms connections lazily via ServerMcpPool.
|
|
|
|
:param spec: Ignored.
|
|
"""
|
|
del spec
|
|
|
|
async def shutdown(self) -> None:
|
|
"""No-op — the Omnigent server owns and manages MCP connections."""
|