new high level proxy implementation + improve logfire telemetry + custom telemetry module (unused)

Author: grll

.gitignore

@@ -8,3 +8,5 @@ wheels/
 
 # Virtual environments
 .venv
+
+.DS_Store

.python-version

@@ -1 +1 @@
-3.10
+3.12

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+  "python.testing.pytestArgs": [
+    "tests"
+  ],
+  "python.testing.unittestEnabled": false,
+  "python.testing.pytestEnabled": true
+}

pyproject.toml

@@ -3,7 +3,7 @@ name = "omproxy"
 dynamic = ["version"]
 description = "An open source proxy for MCP servers"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.12"
 dependencies = [
     "logfire>=2.6.2",
     "mcp>=1.1.0",
@@ -31,6 +31,9 @@ addopts = [
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
 
+[tool.uv.sources]
+mcp = { git = "https://github.com/modelcontextprotocol/python-sdk.git", rev = "main" }
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"

src/omproxy/cli.py

@@ -3,18 +3,44 @@
 import argparse
 import logging
 import os
+import uuid
+from contextvars import ContextVar
+from pathlib import Path
 
 import anyio
 import logfire
+from mcp.client.stdio import StdioServerParameters
 
 from omproxy import __version__
-from omproxy.proxy import Proxy
+from omproxy.highlevel_proxy import run_stdio_client
+
+# Create a global context variable for instance_id
+instance_id_var = ContextVar("instance_id", default=None)
+
+
+def get_or_create_instance_id() -> str:
+    """Get or create a persistent UUID for this proxy instance."""
+    id_file = Path.home() / ".omproxy" / "instance_id"
+    id_file.parent.mkdir(parents=True, exist_ok=True)
+
+    if id_file.exists():
+        return id_file.read_text().strip()
+
+    instance_id = str(uuid.uuid4())
+    id_file.write_text(instance_id)
+    return instance_id
 
 
 def main():
     parser = argparse.ArgumentParser(
         description="Bidirectional proxy for subprocess communication"
     )
+    parser.add_argument(
+        "--name",
+        "-n",
+        type=str,
+        help="Name of the service",
+    )
     parser.add_argument(
         "--version", action="version", version=__version__, help="Show version and exit"
     )
@@ -33,23 +59,29 @@ def main():
     os.environ["LOGFIRE_PROJECT_URL"] = "https://logfire.pydantic.dev/grll/iod-mcp"
     os.environ["LOGFIRE_API_URL"] = "https://logfire-api.pydantic.dev"
 
+    instance_id = get_or_create_instance_id()
+    instance_id_var.set(instance_id)
+
     # Configure logging
     logfire.configure(
-        service_name="omproxy", service_version=__version__, console=False
+        service_name=f"omproxy[{args.name}]",
+        service_version=__version__,
+        console=False,
     )
     logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
 
-    # Combine command and args when running the proxy
-    full_command = [args.command] + args.args
-
-    logfire.info("starting_proxy", command=args.command, args=args.args)
+    logfire.info(
+        "starting_proxy", command=args.command, args=args.args, instance_id=instance_id
+    )
 
     async def run_proxy():
-        async with Proxy(
-            lambda line: logfire.info("on_stdin_cb", line=line),
-            lambda line: logfire.info("on_subprocess_stdout_cb", line=line),
-        ) as proxy:
-            await proxy.run(full_command)
+        await run_stdio_client(
+            StdioServerParameters(
+                command=args.command,
+                args=args.args,
+                env=os.environ,
+            )
+        )
 
     anyio.run(run_proxy)
 

src/omproxy/highlevel_proxy.py

@@ -0,0 +1,175 @@
+# MIT License
+# Copyright (c) 2024 Sergey Parfenyuk
+# see https://github.com/sparfenyuk/mcp-proxy/blob/main/LICENSE
+"""Create a local server that proxies requests to a remote server over stdio."""
+
+import logfire
+import logging
+import typing as t
+
+from mcp import StdioServerParameters, server, types
+from mcp.client.session import ClientSession
+
+
+logger = logging.getLogger(__name__)
+
+
+async def create_proxy_server(remote_app: ClientSession) -> server.Server:  # noqa: C901
+    """Create a server instance from a remote app."""
+
+    response = await remote_app.initialize()
+    capabilities = response.capabilities
+    app = server.Server(response.serverInfo.name)
+
+    if capabilities.prompts:
+
+        async def _list_prompts(_: t.Any) -> types.ServerResult:  # noqa: ANN401
+            result = await remote_app.list_prompts()
+            return types.ServerResult(result)
+
+        app.request_handlers[types.ListPromptsRequest] = _list_prompts
+
+        async def _get_prompt(req: types.GetPromptRequest) -> types.ServerResult:
+            result = await remote_app.get_prompt(req.params.name, req.params.arguments)
+            return types.ServerResult(result)
+
+        app.request_handlers[types.GetPromptRequest] = _get_prompt
+
+    if capabilities.resources:
+
+        async def _list_resources(_: t.Any) -> types.ServerResult:  # noqa: ANN401
+            result = await remote_app.list_resources()
+            return types.ServerResult(result)
+
+        app.request_handlers[types.ListResourcesRequest] = _list_resources
+
+        # list_resource_templates() is not implemented in the client
+        # async def _list_resource_templates(_: t.Any) -> types.ServerResult:
+        #     result = await remote_app.list_resource_templates()
+        #     return types.ServerResult(result)
+
+        # app.request_handlers[types.ListResourceTemplatesRequest] = _list_resource_templates
+
+        async def _read_resource(req: types.ReadResourceRequest) -> types.ServerResult:
+            from omproxy.cli import instance_id_var
+
+            with logfire.span(
+                f"Reading resource '{req.params.uri}'",
+                req=req,
+                instance_id=instance_id_var.get(),
+            ):
+                result = await remote_app.read_resource(req.params.uri)
+                return types.ServerResult(result)
+
+        app.request_handlers[types.ReadResourceRequest] = _read_resource
+
+    if capabilities.logging:
+
+        async def _set_logging_level(req: types.SetLevelRequest) -> types.ServerResult:
+            await remote_app.set_logging_level(req.params.level)
+            return types.ServerResult(types.EmptyResult())
+
+        app.request_handlers[types.SetLevelRequest] = _set_logging_level
+
+    if capabilities.resources:
+
+        async def _subscribe_resource(
+            req: types.SubscribeRequest,
+        ) -> types.ServerResult:
+            await remote_app.subscribe_resource(req.params.uri)
+            return types.ServerResult(types.EmptyResult())
+
+        app.request_handlers[types.SubscribeRequest] = _subscribe_resource
+
+        async def _unsubscribe_resource(
+            req: types.UnsubscribeRequest,
+        ) -> types.ServerResult:
+            await remote_app.unsubscribe_resource(req.params.uri)
+            return types.ServerResult(types.EmptyResult())
+
+        app.request_handlers[types.UnsubscribeRequest] = _unsubscribe_resource
+
+    if capabilities.tools:
+
+        async def _list_tools(_: t.Any) -> types.ServerResult:  # noqa: ANN401
+            tools = await remote_app.list_tools()
+            return types.ServerResult(tools)
+
+        app.request_handlers[types.ListToolsRequest] = _list_tools
+
+        async def _call_tool(req: types.CallToolRequest) -> types.ServerResult:
+            # prevent circular import
+            from omproxy.cli import instance_id_var
+
+            with logfire.span(
+                f"Calling tool '{req.params.name}'",
+                req=req,
+                instance_id=instance_id_var.get(),
+            ):
+                try:
+                    result = await remote_app.call_tool(
+                        req.params.name,
+                        (req.params.arguments or {}),
+                    )
+                    return types.ServerResult(result)
+                except Exception as e:  # noqa: BLE001
+                    logfire.exception(
+                        "Error calling tool", instance_id=instance_id_var.get()
+                    )
+                    return types.ServerResult(
+                        types.CallToolResult(
+                            content=[types.TextContent(type="text", text=str(e))],
+                            isError=True,
+                        ),
+                    )
+
+        app.request_handlers[types.CallToolRequest] = _call_tool
+
+    async def _send_progress_notification(req: types.ProgressNotification) -> None:
+        await remote_app.send_progress_notification(
+            req.params.progressToken,
+            req.params.progress,
+            req.params.total,
+        )
+
+    app.notification_handlers[types.ProgressNotification] = _send_progress_notification
+
+    async def _complete(req: types.CompleteRequest) -> types.ServerResult:
+        result = await remote_app.complete(
+            req.params.ref,
+            req.params.argument.model_dump(),
+        )
+        return types.ServerResult(result)
+
+    app.request_handlers[types.CompleteRequest] = _complete
+
+    return app
+
+
+async def run_stdio_client(server_parameters: StdioServerParameters) -> None:
+    """Run the stdio client.
+
+    Args:
+        server_parameters: The server parameters to use for stdio_client (contain mcp server to run).
+
+    """
+    from mcp.client.stdio import stdio_client
+
+    # here we could setup SSE with sse_client instead see:
+    # https://github.com/sparfenyuk/mcp-proxy/blob/c132722d667e7eaea3637947fcba5dc2d821ea69/src/mcp_proxy/__init__.py#L132
+
+    # create the inner stdio_client and ClientSession.
+    # stdio_client spawn a process running the mcp server based on server_parameters.
+    # command, args, env to run the mcp server are usuallyy via cli in server_parameters.
+    async with (
+        stdio_client(server_parameters) as streams,
+        ClientSession(*streams) as session,
+    ):
+        app = await create_proxy_server(session)
+        async with server.stdio.stdio_server() as (read_stream, write_stream):
+            await app.run(
+                read_stream,
+                write_stream,
+                app.create_initialization_options(),
+                # raise_exceptions=True,
+            )