bohr-agent-sdk 0.1.112__py3-none-any.whl → 0.1.113__py3-none-any.whl

{bohr_agent_sdk-0.1.112.dist-info → bohr_agent_sdk-0.1.113.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bohr-agent-sdk
-Version: 0.1.112
+Version: 0.1.113
 Summary: SDK for scientific agents
 Home-page: https://github.com/dptech-corp/bohr-agent-sdk/
 Author: DP Technology

{bohr_agent_sdk-0.1.112.dist-info → bohr_agent_sdk-0.1.113.dist-info}/RECORD

@@ -62,7 +62,7 @@ dp/agent/device/device/__init__.py,sha256=w7_1S16S1vWUq0RGl0GFgjq2vFkc5oNvy8cQTn
 dp/agent/device/device/device.py,sha256=9ZRIJth-4qMO-i-u_b_cO3d6a4eTbTQjPaxFsV_zEkc,9643
 dp/agent/device/device/types.py,sha256=JuxB-hjf1CjjvfBxCLwRAXVFlYS-nPEdiJpBWLFVCzo,1924
 dp/agent/server/__init__.py,sha256=rckaYd8pbYyB4ENEhgjXKeGMXjdnrgcJpdM1gu5u1Wc,508
-dp/agent/server/calculation_mcp_server.py,sha256=dMIFUQttO7yZyivcuf4BWEZra0Fkh7M7fB2b8tpFix0,17203
+dp/agent/server/calculation_mcp_server.py,sha256=a0hKNVz-WoUbL8y9GhDx1hO830frvmdtvXTBf7V40lI,18478
 dp/agent/server/preprocessor.py,sha256=XUWu7QOwo_sIDMYS2b1OTrM33EXEVH_73vk-ju1Ok8A,1264
 dp/agent/server/utils.py,sha256=cIKaAg8UaP5yMwvIVTgUVBjy-B3S16bEdnucUf4UDIM,2055
 dp/agent/server/executor/__init__.py,sha256=s95M5qKQk39Yi9qaVJZhk_nfj54quSf7EDghR3OCFUA,248
@@ -75,8 +75,8 @@ dp/agent/server/storage/bohrium_storage.py,sha256=EsKX4dWWvZTn2TEhZv4zsvihfDK0mm
 dp/agent/server/storage/http_storage.py,sha256=KiySq7g9-iJr12XQCKKyJLn8wJoDnSRpQAR5_qPJ1ZU,1471
 dp/agent/server/storage/local_storage.py,sha256=t1wfjByjXew9ws3PuUxWxmZQ0-Wt1a6t4wmj3fW62GI,1352
 dp/agent/server/storage/oss_storage.py,sha256=pgjmi7Gir3Y5wkMDCvU4fvSls15fXT7Ax-h9MYHFPK0,3359
-bohr_agent_sdk-0.1.112.dist-info/METADATA,sha256=qnbHAl8Ey0kVMpxD_3Q6qWhOZCC0uocIZXwUQZk11_s,11070
-bohr_agent_sdk-0.1.112.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bohr_agent_sdk-0.1.112.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
-bohr_agent_sdk-0.1.112.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
-bohr_agent_sdk-0.1.112.dist-info/RECORD,,
+bohr_agent_sdk-0.1.113.dist-info/METADATA,sha256=66soACYq7kToAXVjiDX4A_jRkddvoSOLah9QPFrUBxs,11070
+bohr_agent_sdk-0.1.113.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.113.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.113.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.113.dist-info/RECORD,,

dp/agent/server/calculation_mcp_server.py

@@ -17,6 +17,7 @@ from mcp.server.fastmcp.utilities.func_metadata import (
     ArgModelBase,
     func_metadata,
 )
+from mcp.server.sse import SseServerTransport
 from pydantic import BaseModel, Field, create_model
 from starlette.responses import JSONResponse
 from starlette.routing import Route
@@ -267,9 +268,35 @@ class SubmitResult(BaseModel):
     extra_info: dict | None = None
 
 
+def patch_mcp_close_connection():
+    _mock_orig_handle_post_message = SseServerTransport.handle_post_message
+
+    async def _mock_handle_post_message_with_close(self, scope, receive, send):
+        async def _send(message):
+            if message.get("type") == "http.response.start":
+                headers = list(message.get("headers", []))
+                headers = [
+                    (name, value)
+                    for name, value in headers
+                    if name.lower() != b"connection"
+                ]
+                headers.append((b"connection", b"close"))
+                message["headers"] = headers
+            elif message.get("type") == "http.response.body":
+                message["more_body"] = False
+            await send(message)
+        await _mock_orig_handle_post_message(self, scope, receive, _send)
+
+    if not getattr(SseServerTransport.handle_post_message,
+                   "__patched_close__", False):
+        SseServerTransport.handle_post_message = \
+            _mock_handle_post_message_with_close
+        SseServerTransport.handle_post_message.__patched_close__ = True
+
+
 class CalculationMCPServer:
     def __init__(self, *args, preprocess_func=None, fastmcp_mode=False,
-                 **kwargs):
+                 patch_close_connection=False, **kwargs):
         """
         Args:
             preprocess_func: The preprocess function for all tools
@@ -277,6 +304,8 @@ class CalculationMCPServer:
         """
         self.preprocess_func = preprocess_func
         self.fastmcp_mode = fastmcp_mode
+        if patch_close_connection:
+            patch_mcp_close_connection()
         self.mcp = FastMCP(*args, **kwargs)
         self.fn_metadata_map = {}