From 8070681a385bf688c73f3e932e1050461ca47947 Mon Sep 17 00:00:00 2001
From: neatguycoding <15627489+NeatGuyCoding@users.noreply.github.com>
Date: Thu, 10 Jul 2025 19:17:23 +0800
Subject: [PATCH] fix: Potential ThreadPool Leakage in base_session

Signed-off-by: neatguycoding <15627489+NeatGuyCoding@users.noreply.github.com>
---
 api/core/mcp/session/base_session.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/api/core/mcp/session/base_session.py b/api/core/mcp/session/base_session.py
index dd2ab25527..42aeea3908 100644
--- a/api/core/mcp/session/base_session.py
+++ b/api/core/mcp/session/base_session.py
@@ -1,7 +1,7 @@
 import logging
 import queue
 from collections.abc import Callable
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import Future, ThreadPoolExecutor
 from contextlib import ExitStack
 from datetime import timedelta
 from types import TracebackType
@@ -172,8 +172,8 @@ class BaseSession(
         self._in_flight = {}
         self._exit_stack = ExitStack()
         # Initialize executor and future to None for proper cleanup checks
-        self._executor = None
-        self._receiver_future = None
+        self._executor: ThreadPoolExecutor | None = None
+        self._receiver_future: Future | None = None
 
     def __enter__(self) -> Self:
         self._executor = ThreadPoolExecutor()
@@ -181,7 +181,7 @@ class BaseSession(
         return self
 
     def check_receiver_status(self) -> None:
-        if self._receiver_future.done():
+        if self._receiver_future and self._receiver_future.done():
             self._receiver_future.result()
 
     def __exit__(
@@ -200,7 +200,7 @@ class BaseSession(
 
         # Shutdown the executor
         if self._executor:
-            self._executor.shutdown(wait=True, timeout=5.0)
+            self._executor.shutdown(wait=True)
 
     def send_request(
         self,