blaxel 0.2.0rc2__py3-none-any.whl → 0.2.0rc3__py3-none-any.whl

blaxel/core/mcp/client.py

@@ -7,6 +7,7 @@ import anyio
 import mcp.types as types
 from anyio.abc import TaskStatus
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp.shared.message import SessionMessage
 from websockets.asyncio.client import ClientConnection
 from websockets.asyncio.client import connect as ws_connect
 
@@ -28,11 +29,11 @@ async def websocket_client(
 
     The `timeout` parameter controls connection timeout.
     """
-    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
-    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+    read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
 
-    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
-    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+    write_stream: MemoryObjectSendStream[SessionMessage]
+    write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
 
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -64,7 +65,9 @@ async def websocket_client(
                                         message
                                     )
                                     logger.debug(f"Received server message: {parsed_message}")
-                                    await read_stream_writer.send(parsed_message)
+                                    await read_stream_writer.send(
+                                        SessionMessage(message=parsed_message)
+                                    )
                                 except Exception as exc:
                                     logger.error(f"Error parsing server message: {exc}")
                                     await read_stream_writer.send(exc)
@@ -77,10 +80,12 @@ async def websocket_client(
                     async def ws_writer(websocket: ClientConnection):
                         try:
                             async with write_stream_reader:
-                                async for message in write_stream_reader:
-                                    logger.debug(f"Sending client message: {message}")
+                                async for session_message in write_stream_reader:
+                                    logger.debug(
+                                        f"Sending client message: {session_message.message}"
+                                    )
                                     await websocket.send(
-                                        message.model_dump_json(
+                                        session_message.message.model_dump_json(
                                             by_alias=True,
                                             exclude_none=True,
                                         )

blaxel/core/mcp/server.py

@@ -8,6 +8,7 @@ import anyio
 import mcp.types as types
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from mcp.server.fastmcp import FastMCP as FastMCPBase
+from mcp.shared.message import SessionMessage
 from websockets.asyncio.server import ServerConnection, serve
 
 from ..common.env import env
@@ -34,11 +35,11 @@ class BlaxelMcpServerTransport:
     @asynccontextmanager
     async def websocket_server(self):
         """Create and run a WebSocket server for MCP communication."""
-        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
-        read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+        read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
+        read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
 
-        write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
-        write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+        write_stream: MemoryObjectSendStream[SessionMessage]
+        write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
 
         read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
         write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -56,7 +57,7 @@ class BlaxelMcpServerTransport:
                         if hasattr(msg, "id") and msg.id is not None:
                             original_id = msg.id
                             msg.id = f"{client_id}:{original_id}"
-                        await read_stream_writer.send(msg)
+                        await read_stream_writer.send(SessionMessage(message=msg))
                     except Exception as exc:
                         logger.error(f"Failed to parse message: {exc}\n{traceback.format_exc()}")
                         await read_stream_writer.send(exc)
@@ -69,11 +70,11 @@ class BlaxelMcpServerTransport:
 
         async def message_sender():
             async with write_stream_reader:
-                async for message in write_stream_reader:
+                async for session_message in write_stream_reader:
                     # Extract client ID from message ID
                     client_id = None
                     msg_id = None
-
+                    message = session_message.message
                     if hasattr(message, "id") and message.id is not None:
                         parts = str(message.id).split(":", 1)
                         if len(parts) == 2:

{blaxel-0.2.0rc2.dist-info → blaxel-0.2.0rc3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blaxel
-Version: 0.2.0rc2
+Version: 0.2.0rc3
 Summary: Blaxel - AI development platform SDK
 Project-URL: Homepage, https://blaxel.ai
 Project-URL: Documentation, https://docs.blaxel.ai
@@ -11,7 +11,7 @@ License-File: LICENSE
 Requires-Python: >=3.10
 Requires-Dist: attrs>=21.3.0
 Requires-Dist: httpx>=0.27.0
-Requires-Dist: mcp<=1.7.1
+Requires-Dist: mcp>=1.9.4
 Requires-Dist: pydantic>=2.0.0
 Requires-Dist: pyjwt>=2.0.0
 Requires-Dist: python-dateutil>=2.8.0

{blaxel-0.2.0rc2.dist-info → blaxel-0.2.0rc3.dist-info}/RECORD

@@ -301,8 +301,8 @@ blaxel/core/common/logger.py,sha256=Jt0MCJgYDPq36rl7UyKRDJH76a-AwYdfggNeNYJt6N0,
 blaxel/core/common/settings.py,sha256=7KTryuBdud0IfHqykX7xEEtpgq5M5h1Z8YEzYKsHB-Q,2327
 blaxel/core/jobs/__init__.py,sha256=tApZ6t6Savd5c9yH1RhuiPDGq19B464Fakd-YqCOEtI,5441
 blaxel/core/mcp/__init__.py,sha256=5VjkiQFb1QWW5QKRgwPHARlxZJ9Xqaz0diJTpM8LLF0,142
-blaxel/core/mcp/client.py,sha256=dFjPHirybqA3IMecSvdDss-iU5ferPUaunNrPyhFm8o,5195
-blaxel/core/mcp/server.py,sha256=lx-4C71PdPxbG_tGuk81j44jrl7a1ztSixC-k4ObyBo,5812
+blaxel/core/mcp/client.py,sha256=aK3wSnsO8DmT1BZqw4eiCMF71Jwvni6Qga0DhPP806Y,5437
+blaxel/core/mcp/server.py,sha256=tXySGZKgK3IllYOzYOecp58BixKBkmAIvQp_4nSM_Ww,5919
 blaxel/core/models/__init__.py,sha256=HbRDsMnUFHkPC-MMkzPXh4mUqkVjqO6p3j7m00N_XSo,1722
 blaxel/core/sandbox/__init__.py,sha256=eWTH7PWNQZe9efqParwgXzltc7L4OyPtT_REqCdQybU,551
 blaxel/core/sandbox/action.py,sha256=9Zjkco7YkLzBThD3N2Hr5SpeEiqU_-Ktk8HlKpkpiAg,2802
@@ -402,7 +402,7 @@ blaxel/telemetry/instrumentation/map.py,sha256=EVCw6ug1hemE-9kh77Vxs9cqr_p9TPzFw
 blaxel/telemetry/instrumentation/utils.py,sha256=KInMYZH-mu9_wvetmf0EmgrfN3Sw8IWk2Y95v2u90_U,1901
 blaxel/telemetry/log/log.py,sha256=RvQByRjZMoP_dRaAZu8oK6DTegsHs-xV4W-UIqis6CA,2461
 blaxel/telemetry/log/logger.py,sha256=NPAS3g82ryROjvc_DEZaTIfrcehoLEZoP-JkLxADxc0,4113
-blaxel-0.2.0rc2.dist-info/METADATA,sha256=g08ZNxaAovSAaeMwRLMFeoY5YJ2HWAOmTL1e-QqIHko,9878
-blaxel-0.2.0rc2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-blaxel-0.2.0rc2.dist-info/licenses/LICENSE,sha256=p5PNQvpvyDT_0aYBDgmV1fFI_vAD2aSV0wWG7VTgRis,1069
-blaxel-0.2.0rc2.dist-info/RECORD,,
+blaxel-0.2.0rc3.dist-info/METADATA,sha256=6uAwLfyFgd98jymj7QVZMz6et4U1oAr3JgqyL0ZcMkQ,9878
+blaxel-0.2.0rc3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+blaxel-0.2.0rc3.dist-info/licenses/LICENSE,sha256=p5PNQvpvyDT_0aYBDgmV1fFI_vAD2aSV0wWG7VTgRis,1069
+blaxel-0.2.0rc3.dist-info/RECORD,,