make transport functional

Author: SecretiveShell

mcp_bridge/mcp_clients/transports/docker.py

@@ -1,13 +1,13 @@
-import docker.client
+from aiodocker import Docker
 from contextlib import asynccontextmanager
 import anyio
 import anyio.lowlevel
-import anyio.to_thread
 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):
     """
@@ -23,66 +23,72 @@ async def docker_client(server: DockerMCPServer):
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
 
-    client = docker.client.from_env()
-    container = client.containers.run(
-        image=server.image,
-        name=server.container_name if server.container_name else None,
-        command=server.args if server.args else None,
-        environment=server.env if isinstance(server.env, (dict, list)) else {},
-        detach=True,
-        stdin_open=True,
-        stdout=True,
-        stderr=True,
-        tty=False,
-    )
+    docker = Docker()
 
-    logger.debug(f"made instance of docker client for {container.name}")
+    try:
+        # Pull the image if not available
+        await docker.images.pull(server.image)
 
-    # Attach to container's input/output streams
-    attach_socket = container.attach_socket(params={"stdin": 1, "stdout": 1, "stderr": 1, "stream": 1})
-    attach_socket._sock.setblocking(False)
+        # 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},
+        })
 
-    async def read_from_stdout():
-        try:
-            async with read_stream_writer:
-                buffer = ""
-                while True:
-                    chunk = await anyio.to_thread.run_sync(lambda: attach_socket._sock.recv(1024))
-                    if not chunk:
-                        break
+        await container.start()
+        logger.debug(f"Started Docker container {container.id}")
 
-                    chunk = chunk.decode("utf-8")
-                    lines = (buffer + chunk).split("\n")
-                    buffer = lines.pop()
+        # Attach to the container's input/output streams
+        attach_result = container.attach(stdout=True, stdin=True)
 
-                    for line in lines:
-                        try:
-                            message = types.JSONRPCMessage.model_validate_json(line)
-                        except Exception as exc:
-                            await read_stream_writer.send(exc)
+        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()
 
-                        await read_stream_writer.send(message)
-        except anyio.ClosedResourceError:
-            await anyio.lowlevel.checkpoint()
+                            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()
 
-    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 anyio.to_thread.run_sync(lambda: attach_socket._sock.sendall((json + "\n").encode("utf-8")))
-        except anyio.ClosedResourceError:
-            await anyio.lowlevel.checkpoint()
+            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()
 
-    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
     except Exception as e:
         logger.error(f"Error in docker client: {e}")
     finally:
-        attach_socket._sock.close()
-        # Do not stop or remove the container
-        logger.debug(f"Container {container.name} remains running.")
+        await docker.close()
+        logger.debug("Docker client closed.")

pyproject.toml

@@ -5,8 +5,8 @@ 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",
-    "docker>=7.1.0",
     "fastapi>=0.115.6",
     "httpx>=0.28.1",
     "httpx-sse>=0.4.0",