Spaces:
Running
Running
jfouret commited on
Commit ·
6c9d31c
1
Parent(s): 87a2219
add init_timeout for mcp client
Browse files- src/fastmcp/client/client.py +15 -1
src/fastmcp/client/client.py
CHANGED
|
@@ -62,6 +62,8 @@ class Client:
|
|
| 62 |
message_handler: Optional handler for protocol messages
|
| 63 |
progress_handler: Optional handler for progress notifications
|
| 64 |
timeout: Optional timeout for requests (seconds or timedelta)
|
|
|
|
|
|
|
| 65 |
|
| 66 |
Examples:
|
| 67 |
```python
|
|
@@ -93,6 +95,7 @@ class Client:
|
|
| 93 |
message_handler: MessageHandler | None = None,
|
| 94 |
progress_handler: ProgressHandler | None = None,
|
| 95 |
timeout: datetime.timedelta | float | int | None = None,
|
|
|
|
| 96 |
):
|
| 97 |
self.transport = infer_transport(transport)
|
| 98 |
self._session: ClientSession | None = None
|
|
@@ -110,6 +113,17 @@ class Client:
|
|
| 110 |
|
| 111 |
if isinstance(timeout, int | float):
|
| 112 |
timeout = datetime.timedelta(seconds=timeout)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 113 |
|
| 114 |
self._session_kwargs: SessionKwargs = {
|
| 115 |
"sampling_callback": None,
|
|
@@ -168,7 +182,7 @@ class Client:
|
|
| 168 |
self._session = session
|
| 169 |
# Initialize the session
|
| 170 |
try:
|
| 171 |
-
with anyio.fail_after(
|
| 172 |
self._initialize_result = await self._session.initialize()
|
| 173 |
yield
|
| 174 |
except TimeoutError:
|
|
|
|
| 62 |
message_handler: Optional handler for protocol messages
|
| 63 |
progress_handler: Optional handler for progress notifications
|
| 64 |
timeout: Optional timeout for requests (seconds or timedelta)
|
| 65 |
+
init_timeout: Optional timeout for initial connection (seconds or
|
| 66 |
+
timedelta)
|
| 67 |
|
| 68 |
Examples:
|
| 69 |
```python
|
|
|
|
| 95 |
message_handler: MessageHandler | None = None,
|
| 96 |
progress_handler: ProgressHandler | None = None,
|
| 97 |
timeout: datetime.timedelta | float | int | None = None,
|
| 98 |
+
init_timeout: datetime.timedelta | float | int | None = 1,
|
| 99 |
):
|
| 100 |
self.transport = infer_transport(transport)
|
| 101 |
self._session: ClientSession | None = None
|
|
|
|
| 113 |
|
| 114 |
if isinstance(timeout, int | float):
|
| 115 |
timeout = datetime.timedelta(seconds=timeout)
|
| 116 |
+
|
| 117 |
+
if isinstance(init_timeout, int):
|
| 118 |
+
self._init_timeout = float(init_timeout)
|
| 119 |
+
elif isinstance(init_timeout, datetime.timedelta):
|
| 120 |
+
self._init_timeout = float(init_timeout.total_seconds())
|
| 121 |
+
elif isinstance(init_timeout, float):
|
| 122 |
+
self._init_timeout = init_timeout
|
| 123 |
+
else:
|
| 124 |
+
raise ValueError(
|
| 125 |
+
"init_timeout must be int, float or datetime.timedelta"
|
| 126 |
+
)
|
| 127 |
|
| 128 |
self._session_kwargs: SessionKwargs = {
|
| 129 |
"sampling_callback": None,
|
|
|
|
| 182 |
self._session = session
|
| 183 |
# Initialize the session
|
| 184 |
try:
|
| 185 |
+
with anyio.fail_after(self._init_timeout):
|
| 186 |
self._initialize_result = await self._session.initialize()
|
| 187 |
yield
|
| 188 |
except TimeoutError:
|