jfouret commited on
Commit
6c9d31c
·
1 Parent(s): 87a2219

add init_timeout for mcp client

Browse files
Files changed (1) hide show
  1. 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(1):
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: