Merge pull request #25 from SecretiveShell/docker-client

Add docker client support

Author: SecretiveShell

mcp_bridge/config/final.py

@@ -24,9 +24,15 @@ class SSEMCPServer(BaseModel):
     # TODO: expand this once I find a good definition for this
     url: str = Field(description="URL of the MCP server")
 
+class DockerMCPServer(BaseModel):
+    container_name: str | None = Field(default=None, description="Name of the docker container")
+    image: str = Field(description="Image of the docker container")
+    args: list[str] = Field(default_factory=list, description="Command line arguments for the docker container")
+    env: dict[str, str] = Field(default_factory=dict, description="Environment variables for the docker container")
+
 
 MCPServer = Annotated[
-    Union[StdioServerParameters, SSEMCPServer],
+    Union[StdioServerParameters, SSEMCPServer, DockerMCPServer],
     Field(description="MCP server configuration"),
 ]
 

mcp_bridge/mcp_clients/DockerClient.py

@@ -0,0 +1,39 @@
+import asyncio
+from mcp import ClientSession
+from .transports.docker import docker_client
+from config import config
+from config.final import DockerMCPServer
+from .AbstractClient import GenericMcpClient
+from loguru import logger
+
+
+class DockerClient(GenericMcpClient):
+    config: DockerMCPServer
+    session: ClientSession | None = None
+
+    def __init__(self, name: str, config: DockerMCPServer) -> None:
+        super().__init__(name=name)
+
+        self.config = config
+
+    async def _maintain_session(self):
+        async with docker_client(self.config) as client:
+            logger.debug(f"made instance of docker client for {self.name}")
+            async with ClientSession(*client) as session:
+                await session.initialize()
+                logger.debug(f"finished initialise session for {self.name}")
+                self.session = session
+
+                try:
+                    while True:
+                        await asyncio.sleep(10)
+                        if config.logging.log_server_pings:
+                            logger.debug(f"pinging session for {self.name}")
+
+                        await session.send_ping()
+
+                except Exception as exc:
+                    logger.error(f"ping failed for {self.name}: {exc}")
+                    self.session = None
+
+        logger.debug(f"exiting session for {self.name}")

mcp_bridge/mcp_clients/McpClientManager.py

@@ -5,9 +5,10 @@
 
 from .StdioClient import StdioClient
 from .SseClient import SseClient
-from config.final import SSEMCPServer
+from .DockerClient import DockerClient
+from config.final import DockerMCPServer, SSEMCPServer
 
-client_types = Union[StdioClient, SseClient]
+client_types = Union[StdioClient, SseClient, DockerClient]
 
 
 class MCPClientManager:
@@ -36,6 +37,11 @@ async def construct_client(self, name, server_config) -> client_types:
             client = SseClient(name, server_config)  # type: ignore
             await client.start()
             return client
+        
+        if isinstance(server_config, DockerMCPServer):
+            client = DockerClient(name, server_config)
+            await client.start()
+            return client
 
         raise NotImplementedError("Client Type not supported")
 

mcp_bridge/mcp_clients/transports/docker.py

@@ -0,0 +1,94 @@
+from aiodocker import Docker
+from contextlib import asynccontextmanager
+import anyio
+import anyio.lowlevel
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from loguru import logger
+from config.final import DockerMCPServer
+from mcp import types
+
+
+@asynccontextmanager
+async def docker_client(server: DockerMCPServer):
+    """
+    Client transport for Docker: this will connect to a server by
+    running a Docker container and communicating with it over its stdin/stdout.
+    """
+    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+
+    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
+    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+
+    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+    docker = Docker()
+
+    try:
+        # Pull the image if not available
+        await docker.images.pull(server.image)
+
+        # Create the container
+        container = await docker.containers.create({
+            "Image": server.image,
+            "Args": server.args,
+            "OpenStdin": True,
+            "AttachStdout": True,
+            "AttachStderr": True,
+            "Tty": False,
+            "HostConfig": {"AutoRemove": True},
+        })
+
+        await container.start()
+        logger.debug(f"Started Docker container {container.id}")
+
+        # Attach to the container's input/output streams
+        attach_result = container.attach(stdout=True, stdin=True)
+
+        async def read_from_stdout():
+            try:
+                async with read_stream_writer:
+                    buffer = ""
+                    while True: 
+                        msg = await attach_result.read_out()
+                        if msg is None:
+                            continue
+                        chunk = msg.data
+                        if isinstance(chunk, bytes):
+                            chunk = chunk.decode("utf-8")
+                            lines = (buffer + chunk).split("\n")
+                            buffer = lines.pop()
+
+                            for line in lines:
+                                try:
+                                    json_message = types.JSONRPCMessage.model_validate_json(line)
+                                    await read_stream_writer.send(json_message)
+                                except Exception as exc:
+                                    await read_stream_writer.send(exc)
+            except anyio.ClosedResourceError:
+                await anyio.lowlevel.checkpoint()
+
+        async def write_to_stdin():
+            try:
+                async with write_stream_reader:
+                    async for message in write_stream_reader:
+                        json = message.model_dump_json(by_alias=True, exclude_none=True)
+                        await attach_result.write_in(json.encode("utf-8") + b"\n")
+            except anyio.ClosedResourceError:
+                await anyio.lowlevel.checkpoint()
+
+        try:
+            async with anyio.create_task_group() as tg:
+                tg.start_soon(read_from_stdout)
+                tg.start_soon(write_to_stdin)
+                yield read_stream, write_stream
+        finally:
+            await container.stop()
+            await container.delete()
+
+    except Exception as e:
+        logger.error(f"Error in docker client: {e}")
+    finally:
+        await docker.close()
+        logger.debug("Docker client closed.")

pyproject.toml

@@ -5,6 +5,7 @@ description = "A middleware to provide an openAI compatible endpoint that can ca
 readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
+    "aiodocker>=0.24.0",
     "deepmerge>=2.0",
     "fastapi>=0.115.6",
     "httpx>=0.28.1",