avsolatorio commited on
Commit
e6e4f17
·
1 Parent(s): 3f3cebf

Add disconnect on reload

Browse files

Signed-off-by: Aivin V. Solatorio <avsolatorio@gmail.com>

Files changed (1) hide show
  1. mcp_remote_client.py +8 -0
mcp_remote_client.py CHANGED
@@ -134,6 +134,12 @@ class MCPClientWrapper:
134
  self.session = None
135
  return error_msg
136
 
 
 
 
 
 
 
137
  async def process_message(
138
  self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]
139
  ):
@@ -425,6 +431,8 @@ def gradio_interface(
425
  msg.submit(client.process_message, [msg, chatbot], [chatbot, msg])
426
  clear_btn.click(lambda: [], None, chatbot)
427
 
 
 
428
  return demo
429
 
430
 
 
134
  self.session = None
135
  return error_msg
136
 
137
+ async def disconnect(self):
138
+ if self.exit_stack:
139
+ await self.exit_stack.aclose()
140
+ self.exit_stack = None
141
+ self.session = None
142
+
143
  async def process_message(
144
  self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]
145
  ):
 
431
  msg.submit(client.process_message, [msg, chatbot], [chatbot, msg])
432
  clear_btn.click(lambda: [], None, chatbot)
433
 
434
+ demo.unload(client.disconnect)
435
+
436
  return demo
437
 
438