nvidia-nat-mcp 1.3.0rc1__py3-none-any.whl → 1.4.0a20251008__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- nat/meta/pypi.md +2 -2
- nat/plugins/mcp/auth/auth_provider.py +42 -10
- nat/plugins/mcp/auth/auth_provider_config.py +5 -0
- nat/plugins/mcp/auth/register.py +1 -1
- nat/plugins/mcp/auth/token_storage.py +265 -0
- nat/plugins/mcp/client_base.py +103 -111
- nat/plugins/mcp/client_config.py +131 -0
- nat/plugins/mcp/client_impl.py +293 -104
- nat/plugins/mcp/tool.py +5 -0
- nat/plugins/mcp/utils.py +16 -0
- {nvidia_nat_mcp-1.3.0rc1.dist-info → nvidia_nat_mcp-1.4.0a20251008.dist-info}/METADATA +5 -4
- nvidia_nat_mcp-1.4.0a20251008.dist-info/RECORD +21 -0
- nvidia_nat_mcp-1.3.0rc1.dist-info/RECORD +0 -19
- {nvidia_nat_mcp-1.3.0rc1.dist-info → nvidia_nat_mcp-1.4.0a20251008.dist-info}/WHEEL +0 -0
- {nvidia_nat_mcp-1.3.0rc1.dist-info → nvidia_nat_mcp-1.4.0a20251008.dist-info}/entry_points.txt +0 -0
- {nvidia_nat_mcp-1.3.0rc1.dist-info → nvidia_nat_mcp-1.4.0a20251008.dist-info}/top_level.txt +0 -0
nat/plugins/mcp/client_impl.py
CHANGED
@@ -13,29 +13,41 @@
|
|
13
13
|
# See the License for the specific language governing permissions and
|
14
14
|
# limitations under the License.
|
15
15
|
|
16
|
+
import asyncio
|
16
17
|
import logging
|
18
|
+
from contextlib import asynccontextmanager
|
19
|
+
from dataclasses import dataclass
|
20
|
+
from dataclasses import field
|
21
|
+
from datetime import datetime
|
17
22
|
from datetime import timedelta
|
18
|
-
from typing import Literal
|
19
23
|
|
24
|
+
import aiorwlock
|
20
25
|
from pydantic import BaseModel
|
21
|
-
from pydantic import Field
|
22
|
-
from pydantic import HttpUrl
|
23
|
-
from pydantic import model_validator
|
24
26
|
|
27
|
+
from nat.authentication.interfaces import AuthProviderBase
|
25
28
|
from nat.builder.builder import Builder
|
26
29
|
from nat.builder.function import FunctionGroup
|
27
30
|
from nat.cli.register_workflow import register_function_group
|
28
|
-
from nat.
|
29
|
-
from nat.
|
30
|
-
from nat.plugins.mcp.
|
31
|
+
from nat.plugins.mcp.client_base import MCPBaseClient
|
32
|
+
from nat.plugins.mcp.client_config import MCPClientConfig
|
33
|
+
from nat.plugins.mcp.client_config import MCPToolOverrideConfig
|
34
|
+
from nat.plugins.mcp.utils import truncate_session_id
|
31
35
|
|
32
36
|
logger = logging.getLogger(__name__)
|
33
37
|
|
34
38
|
|
39
|
+
@dataclass
|
40
|
+
class SessionData:
|
41
|
+
"""Container for all session-related data."""
|
42
|
+
client: MCPBaseClient
|
43
|
+
last_activity: datetime
|
44
|
+
ref_count: int = 0
|
45
|
+
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
|
46
|
+
|
47
|
+
|
35
48
|
class MCPFunctionGroup(FunctionGroup):
|
36
49
|
"""
|
37
|
-
A specialized FunctionGroup for MCP clients that includes MCP-specific attributes
|
38
|
-
with proper type safety.
|
50
|
+
A specialized FunctionGroup for MCP clients that includes MCP-specific attributes with session management.
|
39
51
|
"""
|
40
52
|
|
41
53
|
def __init__(self, *args, **kwargs):
|
@@ -45,6 +57,20 @@ class MCPFunctionGroup(FunctionGroup):
|
|
45
57
|
self._mcp_client_server_name: str | None = None
|
46
58
|
self._mcp_client_transport: str | None = None
|
47
59
|
|
60
|
+
# Session management - consolidated data structure
|
61
|
+
self._sessions: dict[str, SessionData] = {}
|
62
|
+
|
63
|
+
# Use RWLock for better concurrency: multiple readers (tool calls) can access
|
64
|
+
# existing sessions simultaneously, while writers (create/delete) get exclusive access
|
65
|
+
self._session_rwlock = aiorwlock.RWLock()
|
66
|
+
# Throttled cleanup control
|
67
|
+
self._last_cleanup_check: datetime = datetime.now()
|
68
|
+
self._cleanup_check_interval: timedelta = timedelta(minutes=5)
|
69
|
+
|
70
|
+
# Shared components for session client creation
|
71
|
+
self._shared_auth_provider: AuthProviderBase | None = None
|
72
|
+
self._client_config: MCPClientConfig | None = None
|
73
|
+
|
48
74
|
@property
|
49
75
|
def mcp_client(self):
|
50
76
|
"""Get the MCP client instance."""
|
@@ -75,105 +101,253 @@ class MCPFunctionGroup(FunctionGroup):
|
|
75
101
|
"""Set the MCP client transport type."""
|
76
102
|
self._mcp_client_transport = transport
|
77
103
|
|
104
|
+
@property
|
105
|
+
def session_count(self) -> int:
|
106
|
+
"""Current number of active sessions."""
|
107
|
+
return len(self._sessions)
|
78
108
|
|
79
|
-
|
80
|
-
|
81
|
-
|
82
|
-
|
83
|
-
|
84
|
-
|
85
|
-
|
86
|
-
|
87
|
-
|
88
|
-
|
89
|
-
|
90
|
-
|
91
|
-
|
92
|
-
|
93
|
-
|
94
|
-
|
95
|
-
|
96
|
-
|
97
|
-
|
98
|
-
|
99
|
-
|
100
|
-
|
101
|
-
|
102
|
-
|
103
|
-
|
104
|
-
|
105
|
-
|
106
|
-
|
107
|
-
|
108
|
-
|
109
|
-
|
110
|
-
|
111
|
-
|
112
|
-
|
113
|
-
|
114
|
-
|
115
|
-
|
116
|
-
|
117
|
-
|
118
|
-
|
119
|
-
|
120
|
-
|
121
|
-
|
122
|
-
|
123
|
-
|
124
|
-
|
125
|
-
|
126
|
-
|
127
|
-
|
128
|
-
|
129
|
-
|
130
|
-
|
131
|
-
|
132
|
-
|
133
|
-
|
134
|
-
|
135
|
-
|
136
|
-
|
109
|
+
@property
|
110
|
+
def session_limit(self) -> int:
|
111
|
+
"""Maximum allowed sessions."""
|
112
|
+
return self._client_config.max_sessions if self._client_config else 100
|
113
|
+
|
114
|
+
def _get_session_id_from_context(self) -> str | None:
|
115
|
+
"""Get the session ID from the current context."""
|
116
|
+
try:
|
117
|
+
from nat.builder.context import Context as _Ctx
|
118
|
+
|
119
|
+
# Get session id from context, authentication is done per-websocket session for tool calls
|
120
|
+
session_id = None
|
121
|
+
cookies = getattr(_Ctx.get().metadata, "cookies", None)
|
122
|
+
if cookies:
|
123
|
+
session_id = cookies.get("nat-session")
|
124
|
+
|
125
|
+
if not session_id:
|
126
|
+
# use default user id if allowed
|
127
|
+
if self._shared_auth_provider and \
|
128
|
+
self._shared_auth_provider.config.allow_default_user_id_for_tool_calls:
|
129
|
+
session_id = self._shared_auth_provider.config.default_user_id
|
130
|
+
return session_id
|
131
|
+
except Exception:
|
132
|
+
return None
|
133
|
+
|
134
|
+
async def cleanup_sessions(self, max_age: timedelta | None = None) -> int:
|
135
|
+
"""
|
136
|
+
Manually trigger cleanup of inactive sessions.
|
137
|
+
|
138
|
+
Args:
|
139
|
+
max_age: Maximum age for sessions before cleanup. If None, uses configured timeout.
|
140
|
+
|
141
|
+
Returns:
|
142
|
+
Number of sessions cleaned up.
|
143
|
+
"""
|
144
|
+
sessions_before = len(self._sessions)
|
145
|
+
await self._cleanup_inactive_sessions(max_age)
|
146
|
+
sessions_after = len(self._sessions)
|
147
|
+
return sessions_before - sessions_after
|
148
|
+
|
149
|
+
async def _cleanup_inactive_sessions(self, max_age: timedelta | None = None):
|
150
|
+
"""Remove clients for sessions inactive longer than max_age.
|
151
|
+
|
152
|
+
This method uses the RWLock writer to ensure thread-safe cleanup.
|
153
|
+
"""
|
154
|
+
if max_age is None:
|
155
|
+
max_age = self._client_config.session_idle_timeout if self._client_config else timedelta(hours=1)
|
156
|
+
|
157
|
+
async with self._session_rwlock.writer:
|
158
|
+
current_time = datetime.now()
|
159
|
+
inactive_sessions = []
|
160
|
+
|
161
|
+
for session_id, session_data in self._sessions.items():
|
162
|
+
# Skip cleanup if session is actively being used
|
163
|
+
if session_data.ref_count > 0:
|
164
|
+
continue
|
165
|
+
|
166
|
+
if current_time - session_data.last_activity > max_age:
|
167
|
+
inactive_sessions.append(session_id)
|
168
|
+
|
169
|
+
for session_id in inactive_sessions:
|
170
|
+
try:
|
171
|
+
logger.info("Cleaning up inactive session client: %s", truncate_session_id(session_id))
|
172
|
+
session_data = self._sessions[session_id]
|
173
|
+
# Close the client connection
|
174
|
+
await session_data.client.__aexit__(None, None, None)
|
175
|
+
logger.info("Cleaned up inactive session client: %s", truncate_session_id(session_id))
|
176
|
+
except Exception as e:
|
177
|
+
logger.warning("Error cleaning up session client %s: %s", truncate_session_id(session_id), e)
|
178
|
+
finally:
|
179
|
+
# Always remove from tracking to prevent leaks, even if close failed
|
180
|
+
self._sessions.pop(session_id, None)
|
181
|
+
logger.info("Cleaned up session tracking for: %s", truncate_session_id(session_id))
|
182
|
+
logger.info(" Total sessions: %d", len(self._sessions))
|
183
|
+
|
184
|
+
async def _get_session_client(self, session_id: str) -> MCPBaseClient:
|
185
|
+
"""Get the appropriate MCP client for the session."""
|
186
|
+
# Throttled cleanup on access
|
187
|
+
now = datetime.now()
|
188
|
+
if now - self._last_cleanup_check > self._cleanup_check_interval:
|
189
|
+
await self._cleanup_inactive_sessions()
|
190
|
+
self._last_cleanup_check = now
|
191
|
+
|
192
|
+
# If the session_id equals the configured default_user_id use the base client
|
193
|
+
# instead of creating a per-session client
|
194
|
+
if self._shared_auth_provider:
|
195
|
+
default_uid = self._shared_auth_provider.config.default_user_id
|
196
|
+
if default_uid and session_id == default_uid:
|
197
|
+
return self.mcp_client
|
198
|
+
|
199
|
+
# Fast path: check if session already exists (reader lock for concurrent access)
|
200
|
+
async with self._session_rwlock.reader:
|
201
|
+
if session_id in self._sessions:
|
202
|
+
# Update last activity for existing client
|
203
|
+
self._sessions[session_id].last_activity = datetime.now()
|
204
|
+
return self._sessions[session_id].client
|
205
|
+
|
206
|
+
# Check session limit before creating new client (outside writer lock to avoid deadlock)
|
207
|
+
if self._client_config and len(self._sessions) >= self._client_config.max_sessions:
|
208
|
+
# Try cleanup first to free up space
|
209
|
+
await self._cleanup_inactive_sessions()
|
210
|
+
|
211
|
+
# Slow path: create session with writer lock for exclusive access
|
212
|
+
async with self._session_rwlock.writer:
|
213
|
+
# Double-check after acquiring writer lock (another coroutine might have created it)
|
214
|
+
if session_id in self._sessions:
|
215
|
+
self._sessions[session_id].last_activity = datetime.now()
|
216
|
+
return self._sessions[session_id].client
|
217
|
+
|
218
|
+
# Re-check session limit inside writer lock
|
219
|
+
if self._client_config and len(self._sessions) >= self._client_config.max_sessions:
|
220
|
+
logger.warning("Session limit reached (%d), rejecting new session: %s",
|
221
|
+
self._client_config.max_sessions,
|
222
|
+
truncate_session_id(session_id))
|
223
|
+
raise RuntimeError(f"Service temporarily unavailable: Maximum concurrent sessions "
|
224
|
+
f"({self._client_config.max_sessions}) exceeded. Please try again later.")
|
225
|
+
|
226
|
+
# Create session client lazily
|
227
|
+
logger.info("Creating new MCP client for session: %s", truncate_session_id(session_id))
|
228
|
+
session_client = await self._create_session_client(session_id)
|
229
|
+
|
230
|
+
# Create session data with all components
|
231
|
+
session_data = SessionData(client=session_client, last_activity=datetime.now(), ref_count=0)
|
232
|
+
|
233
|
+
# Cache the session data
|
234
|
+
self._sessions[session_id] = session_data
|
235
|
+
logger.info(" Total sessions: %d", len(self._sessions))
|
236
|
+
return session_client
|
237
|
+
|
238
|
+
@asynccontextmanager
|
239
|
+
async def _session_usage_context(self, session_id: str):
|
240
|
+
"""Context manager to track active session usage and prevent cleanup."""
|
241
|
+
# Ensure session exists - create it if it doesn't
|
242
|
+
if session_id not in self._sessions:
|
243
|
+
# Create session client first
|
244
|
+
await self._get_session_client(session_id)
|
245
|
+
# Session should now exist in _sessions
|
246
|
+
|
247
|
+
# Get session data (session must exist at this point)
|
248
|
+
session_data = self._sessions[session_id]
|
249
|
+
|
250
|
+
# Thread-safe reference counting using per-session lock
|
251
|
+
async with session_data.lock:
|
252
|
+
session_data.ref_count += 1
|
253
|
+
|
254
|
+
try:
|
255
|
+
yield
|
256
|
+
finally:
|
257
|
+
async with session_data.lock:
|
258
|
+
session_data.ref_count -= 1
|
259
|
+
|
260
|
+
async def _create_session_client(self, session_id: str) -> MCPBaseClient:
|
261
|
+
"""Create a new MCP client instance for the session."""
|
262
|
+
from nat.plugins.mcp.client_base import MCPStreamableHTTPClient
|
263
|
+
|
264
|
+
config = self._client_config
|
265
|
+
if not config:
|
266
|
+
raise RuntimeError("Client config not initialized")
|
267
|
+
|
268
|
+
if config.server.transport == "streamable-http":
|
269
|
+
client = MCPStreamableHTTPClient(
|
270
|
+
str(config.server.url),
|
271
|
+
auth_provider=self._shared_auth_provider,
|
272
|
+
user_id=session_id, # Pass session_id as user_id for cache isolation
|
273
|
+
tool_call_timeout=config.tool_call_timeout,
|
274
|
+
auth_flow_timeout=config.auth_flow_timeout,
|
275
|
+
reconnect_enabled=config.reconnect_enabled,
|
276
|
+
reconnect_max_attempts=config.reconnect_max_attempts,
|
277
|
+
reconnect_initial_backoff=config.reconnect_initial_backoff,
|
278
|
+
reconnect_max_backoff=config.reconnect_max_backoff)
|
279
|
+
else:
|
280
|
+
# per-user sessions are only supported for streamable-http transport
|
281
|
+
raise ValueError(f"Unsupported transport: {config.server.transport}")
|
282
|
+
|
283
|
+
# Initialize the client
|
284
|
+
await client.__aenter__()
|
285
|
+
|
286
|
+
logger.info("Created session client for session: %s", truncate_session_id(session_id))
|
287
|
+
return client
|
288
|
+
|
289
|
+
|
290
|
+
def mcp_session_tool_function(tool, function_group: MCPFunctionGroup):
|
291
|
+
"""Create a session-aware NAT function for an MCP tool.
|
292
|
+
|
293
|
+
Routes each invocation to the appropriate per-session MCP client while
|
294
|
+
preserving the original tool input schema, converters, and description.
|
137
295
|
"""
|
138
|
-
|
139
|
-
|
140
|
-
|
141
|
-
|
142
|
-
|
143
|
-
|
144
|
-
|
145
|
-
|
146
|
-
|
147
|
-
|
148
|
-
|
149
|
-
|
150
|
-
|
151
|
-
|
152
|
-
|
153
|
-
|
154
|
-
|
155
|
-
|
156
|
-
|
157
|
-
|
158
|
-
|
159
|
-
|
160
|
-
|
161
|
-
|
162
|
-
|
163
|
-
|
164
|
-
|
165
|
-
|
166
|
-
|
167
|
-
|
168
|
-
|
169
|
-
|
170
|
-
|
296
|
+
from nat.builder.function import FunctionInfo
|
297
|
+
|
298
|
+
def _convert_from_str(input_str: str) -> tool.input_schema:
|
299
|
+
return tool.input_schema.model_validate_json(input_str)
|
300
|
+
|
301
|
+
async def _response_fn(tool_input: BaseModel | None = None, **kwargs) -> str:
|
302
|
+
"""Response function for the session-aware tool."""
|
303
|
+
try:
|
304
|
+
# Route to the appropriate session client
|
305
|
+
session_id = function_group._get_session_id_from_context()
|
306
|
+
|
307
|
+
# If no session is available and default-user fallback is disabled, deny the call
|
308
|
+
if function_group._shared_auth_provider and session_id is None:
|
309
|
+
return "User not authorized to call the tool"
|
310
|
+
|
311
|
+
# Check if this is the default user - if so, use base client directly
|
312
|
+
if (not function_group._shared_auth_provider
|
313
|
+
or session_id == function_group._shared_auth_provider.config.default_user_id):
|
314
|
+
# Use base client directly for default user
|
315
|
+
client = function_group.mcp_client
|
316
|
+
session_tool = await client.get_tool(tool.name)
|
317
|
+
else:
|
318
|
+
# Use session usage context to prevent cleanup during tool execution
|
319
|
+
async with function_group._session_usage_context(session_id):
|
320
|
+
client = await function_group._get_session_client(session_id)
|
321
|
+
session_tool = await client.get_tool(tool.name)
|
322
|
+
|
323
|
+
# Preserve original calling convention
|
324
|
+
if tool_input:
|
325
|
+
args = tool_input.model_dump()
|
326
|
+
return await session_tool.acall(args)
|
327
|
+
|
328
|
+
_ = session_tool.input_schema.model_validate(kwargs)
|
329
|
+
return await session_tool.acall(kwargs)
|
330
|
+
except Exception as e:
|
331
|
+
if tool_input:
|
332
|
+
logger.warning("Error calling tool %s with serialized input: %s",
|
333
|
+
tool.name,
|
334
|
+
tool_input.model_dump(),
|
335
|
+
exc_info=True)
|
336
|
+
else:
|
337
|
+
logger.warning("Error calling tool %s with input: %s", tool.name, kwargs, exc_info=True)
|
338
|
+
return str(e)
|
339
|
+
|
340
|
+
return FunctionInfo.create(single_fn=_response_fn,
|
341
|
+
description=tool.description,
|
342
|
+
input_schema=tool.input_schema,
|
343
|
+
converters=[_convert_from_str])
|
171
344
|
|
172
345
|
|
173
346
|
@register_function_group(config_type=MCPClientConfig)
|
174
347
|
async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
|
175
348
|
"""
|
176
349
|
Connect to an MCP server and expose tools as a function group.
|
350
|
+
|
177
351
|
Args:
|
178
352
|
config: The configuration for the MCP client
|
179
353
|
_builder: The builder
|
@@ -196,7 +370,8 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
|
|
196
370
|
client = MCPStdioClient(config.server.command,
|
197
371
|
config.server.args,
|
198
372
|
config.server.env,
|
199
|
-
config.tool_call_timeout,
|
373
|
+
tool_call_timeout=config.tool_call_timeout,
|
374
|
+
auth_flow_timeout=config.auth_flow_timeout,
|
200
375
|
reconnect_enabled=config.reconnect_enabled,
|
201
376
|
reconnect_max_attempts=config.reconnect_max_attempts,
|
202
377
|
reconnect_initial_backoff=config.reconnect_initial_backoff,
|
@@ -204,14 +379,19 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
|
|
204
379
|
elif config.server.transport == "sse":
|
205
380
|
client = MCPSSEClient(str(config.server.url),
|
206
381
|
tool_call_timeout=config.tool_call_timeout,
|
382
|
+
auth_flow_timeout=config.auth_flow_timeout,
|
207
383
|
reconnect_enabled=config.reconnect_enabled,
|
208
384
|
reconnect_max_attempts=config.reconnect_max_attempts,
|
209
385
|
reconnect_initial_backoff=config.reconnect_initial_backoff,
|
210
386
|
reconnect_max_backoff=config.reconnect_max_backoff)
|
211
387
|
elif config.server.transport == "streamable-http":
|
388
|
+
# Use default_user_id for the base client
|
389
|
+
base_user_id = auth_provider.config.default_user_id if auth_provider else None
|
212
390
|
client = MCPStreamableHTTPClient(str(config.server.url),
|
213
391
|
auth_provider=auth_provider,
|
392
|
+
user_id=base_user_id,
|
214
393
|
tool_call_timeout=config.tool_call_timeout,
|
394
|
+
auth_flow_timeout=config.auth_flow_timeout,
|
215
395
|
reconnect_enabled=config.reconnect_enabled,
|
216
396
|
reconnect_max_attempts=config.reconnect_max_attempts,
|
217
397
|
reconnect_initial_backoff=config.reconnect_initial_backoff,
|
@@ -224,6 +404,10 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
|
|
224
404
|
# Create the MCP function group
|
225
405
|
group = MCPFunctionGroup(config=config)
|
226
406
|
|
407
|
+
# Store shared components for session client creation
|
408
|
+
group._shared_auth_provider = auth_provider
|
409
|
+
group._client_config = config
|
410
|
+
|
227
411
|
async with client:
|
228
412
|
# Expose the live MCP client on the function group instance so other components (e.g., HTTP endpoints)
|
229
413
|
# can reuse the already-established session instead of creating a new client per request.
|
@@ -243,13 +427,17 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
|
|
243
427
|
function_name = override.alias if override and override.alias else tool_name
|
244
428
|
description = override.description if override and override.description else tool.description
|
245
429
|
|
246
|
-
# Create the tool function
|
247
|
-
|
430
|
+
# Create the tool function according to configuration
|
431
|
+
if config.session_aware_tools:
|
432
|
+
tool_fn = mcp_session_tool_function(tool, group)
|
433
|
+
else:
|
434
|
+
from nat.plugins.mcp.tool import mcp_tool_function
|
435
|
+
tool_fn = mcp_tool_function(tool)
|
248
436
|
|
249
437
|
# Normalize optional typing for linter/type-checker compatibility
|
250
438
|
single_fn = tool_fn.single_fn
|
251
439
|
if single_fn is None:
|
252
|
-
# Should not happen because
|
440
|
+
# Should not happen because FunctionInfo always sets a single_fn
|
253
441
|
logger.warning("Skipping tool %s because single_fn is None", function_name)
|
254
442
|
continue
|
255
443
|
|
@@ -273,6 +461,7 @@ def mcp_apply_tool_alias_and_description(
|
|
273
461
|
all_tools: dict, tool_overrides: dict[str, MCPToolOverrideConfig] | None) -> dict[str, MCPToolOverrideConfig]:
|
274
462
|
"""
|
275
463
|
Filter tool overrides to only include tools that exist in the MCP server.
|
464
|
+
|
276
465
|
Args:
|
277
466
|
all_tools: The tools from the MCP server
|
278
467
|
tool_overrides: The tool overrides to apply
|
nat/plugins/mcp/tool.py
CHANGED
@@ -26,6 +26,7 @@ from nat.builder.function_info import FunctionInfo
|
|
26
26
|
from nat.cli.register_workflow import register_function
|
27
27
|
from nat.data_models.function import FunctionBaseConfig
|
28
28
|
from nat.plugins.mcp.client_base import MCPToolClient
|
29
|
+
from nat.utils.decorators import deprecated
|
29
30
|
|
30
31
|
logger = logging.getLogger(__name__)
|
31
32
|
|
@@ -109,6 +110,10 @@ def mcp_tool_function(tool: MCPToolClient) -> FunctionInfo:
|
|
109
110
|
|
110
111
|
|
111
112
|
@register_function(config_type=MCPToolConfig)
|
113
|
+
@deprecated(
|
114
|
+
reason=
|
115
|
+
"This function is being replaced with the new mcp_client function group that supports additional MCP features",
|
116
|
+
feature_name="mcp_tool_wrapper")
|
112
117
|
async def mcp_tool(config: MCPToolConfig, builder: Builder):
|
113
118
|
"""
|
114
119
|
Generate a NeMo Agent Toolkit Function that wraps a tool provided by the MCP server.
|
nat/plugins/mcp/utils.py
CHANGED
@@ -21,6 +21,22 @@ from pydantic import Field
|
|
21
21
|
from pydantic import create_model
|
22
22
|
|
23
23
|
|
24
|
+
def truncate_session_id(session_id: str, max_length: int = 10) -> str:
|
25
|
+
"""
|
26
|
+
Truncate a session ID for logging purposes.
|
27
|
+
|
28
|
+
Args:
|
29
|
+
session_id: The session ID to truncate
|
30
|
+
max_length: Maximum length before truncation (default: 10)
|
31
|
+
|
32
|
+
Returns:
|
33
|
+
Truncated session ID with "..." if longer than max_length, otherwise full ID
|
34
|
+
"""
|
35
|
+
if len(session_id) > max_length:
|
36
|
+
return session_id[:max_length] + "..."
|
37
|
+
return session_id
|
38
|
+
|
39
|
+
|
24
40
|
def model_from_mcp_schema(name: str, mcp_input_schema: dict) -> type[BaseModel]:
|
25
41
|
"""
|
26
42
|
Create a pydantic model from the input schema of the MCP tool
|
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.4
|
2
2
|
Name: nvidia-nat-mcp
|
3
|
-
Version: 1.
|
3
|
+
Version: 1.4.0a20251008
|
4
4
|
Summary: Subpackage for MCP client integration in NeMo Agent toolkit
|
5
5
|
Keywords: ai,rag,agents,mcp
|
6
6
|
Classifier: Programming Language :: Python
|
@@ -9,7 +9,8 @@ Classifier: Programming Language :: Python :: 3.12
|
|
9
9
|
Classifier: Programming Language :: Python :: 3.13
|
10
10
|
Requires-Python: <3.14,>=3.11
|
11
11
|
Description-Content-Type: text/markdown
|
12
|
-
Requires-Dist: nvidia-nat==v1.
|
12
|
+
Requires-Dist: nvidia-nat==v1.4.0a20251008
|
13
|
+
Requires-Dist: aiorwlock~=1.5
|
13
14
|
Requires-Dist: mcp~=1.14
|
14
15
|
|
15
16
|
<!--
|
@@ -33,9 +34,9 @@ limitations under the License.
|
|
33
34
|
|
34
35
|
|
35
36
|
# NVIDIA NeMo Agent Toolkit MCP Subpackage
|
36
|
-
Subpackage for MCP
|
37
|
+
Subpackage for MCP integration in NeMo Agent toolkit.
|
37
38
|
|
38
|
-
This package provides MCP (Model Context Protocol)
|
39
|
+
This package provides MCP (Model Context Protocol) functionality, allowing NeMo Agent toolkit workflows to connect to external MCP servers and use their tools as functions.
|
39
40
|
|
40
41
|
## Features
|
41
42
|
|
@@ -0,0 +1,21 @@
|
|
1
|
+
nat/meta/pypi.md,sha256=EYyJTCCEOWzuuz-uNaYJ_WBk55Jiig87wcUr9E4g0yw,1484
|
2
|
+
nat/plugins/mcp/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
|
3
|
+
nat/plugins/mcp/client_base.py,sha256=nos9NTQ2NlU9vd0PFb1n_q9AZtPKQ5OZ6EU74Ydo1C0,26533
|
4
|
+
nat/plugins/mcp/client_config.py,sha256=l9tVUHe8WdFPJ9rXDg8dZkQi1dvHGYwoqQ8Glqg2LGs,6783
|
5
|
+
nat/plugins/mcp/client_impl.py,sha256=uw_iCOwkwbkHYGRW0XSis3wL3jsmf1RDOO6epVy5UPY,21372
|
6
|
+
nat/plugins/mcp/exception_handler.py,sha256=4JVdZDJL4LyumZEcMIEBK2LYC6djuSMzqUhQDZZ6dUo,7648
|
7
|
+
nat/plugins/mcp/exceptions.py,sha256=EGVOnYlui8xufm8dhJyPL1SUqBLnCGOTvRoeyNcmcWE,5980
|
8
|
+
nat/plugins/mcp/register.py,sha256=HOT2Wl2isGuyFc7BUTi58-BbjI5-EtZMZo7stsv5pN4,831
|
9
|
+
nat/plugins/mcp/tool.py,sha256=TS-64B-vdBrQVTojbAx6k3IlhvN7Gc7E_bxa1Uccc4w,6387
|
10
|
+
nat/plugins/mcp/utils.py,sha256=4kNF5FJRiDUn-3fQcsvwvWtG6tYG1y4jU7vpptp0fsA,4522
|
11
|
+
nat/plugins/mcp/auth/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
|
12
|
+
nat/plugins/mcp/auth/auth_flow_handler.py,sha256=2JgK0aH-5ouQCd2ov0lDMJAD5ZWIQJ7SVcXaLArxn6Y,6010
|
13
|
+
nat/plugins/mcp/auth/auth_provider.py,sha256=BgH66DlZgzhLDLO4cBERpHvNAmli5fMo_SCy11W9aBU,21251
|
14
|
+
nat/plugins/mcp/auth/auth_provider_config.py,sha256=b1AaXzOuAkygKXAWSxMKWg8wfW8k33tmUUq6Dk5Mmwk,4038
|
15
|
+
nat/plugins/mcp/auth/register.py,sha256=L2x69NjJPS4s6CCE5myzWVrWn3e_ttHyojmGXvBipMg,1228
|
16
|
+
nat/plugins/mcp/auth/token_storage.py,sha256=aS13ZvEJXcYzkZ0GSbrSor4i5bpjD5BkXHQw1iywC9k,9240
|
17
|
+
nvidia_nat_mcp-1.4.0a20251008.dist-info/METADATA,sha256=A9ksFEOx30iLCUQw3m2s8KhitBt5hAyKL-oL46NeWn8,2013
|
18
|
+
nvidia_nat_mcp-1.4.0a20251008.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
|
19
|
+
nvidia_nat_mcp-1.4.0a20251008.dist-info/entry_points.txt,sha256=rYvUp4i-klBr3bVNh7zYOPXret704vTjvCk1qd7FooI,97
|
20
|
+
nvidia_nat_mcp-1.4.0a20251008.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
|
21
|
+
nvidia_nat_mcp-1.4.0a20251008.dist-info/RECORD,,
|
@@ -1,19 +0,0 @@
|
|
1
|
-
nat/meta/pypi.md,sha256=GyV4DI1d9ThgEhnYTQ0vh40Q9hPC8jN-goLnRiFDmZ8,1498
|
2
|
-
nat/plugins/mcp/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
|
3
|
-
nat/plugins/mcp/client_base.py,sha256=UPdQ4vxe4MXVTF05jJx20dn8JPTvQN1ugZ-dups3-44,26418
|
4
|
-
nat/plugins/mcp/client_impl.py,sha256=sjSFIJeD6LkexN5IbDq_OuoKW52ulpp26LpwaUjOpwg,13142
|
5
|
-
nat/plugins/mcp/exception_handler.py,sha256=4JVdZDJL4LyumZEcMIEBK2LYC6djuSMzqUhQDZZ6dUo,7648
|
6
|
-
nat/plugins/mcp/exceptions.py,sha256=EGVOnYlui8xufm8dhJyPL1SUqBLnCGOTvRoeyNcmcWE,5980
|
7
|
-
nat/plugins/mcp/register.py,sha256=HOT2Wl2isGuyFc7BUTi58-BbjI5-EtZMZo7stsv5pN4,831
|
8
|
-
nat/plugins/mcp/tool.py,sha256=v3MFsiaLJy8Ourcfqa6ohtAE2Nn-vqpC6Q6gsCdJ28Q,6165
|
9
|
-
nat/plugins/mcp/utils.py,sha256=3fuzYpC14wrfMOTOGvY2KHWcxZvBWqrxdDZD17lhmC8,4055
|
10
|
-
nat/plugins/mcp/auth/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
|
11
|
-
nat/plugins/mcp/auth/auth_flow_handler.py,sha256=2JgK0aH-5ouQCd2ov0lDMJAD5ZWIQJ7SVcXaLArxn6Y,6010
|
12
|
-
nat/plugins/mcp/auth/auth_provider.py,sha256=OfxPCEaXuhP8anOdrTRH-_E78CrbJtzW6i81_kebpDk,19321
|
13
|
-
nat/plugins/mcp/auth/auth_provider_config.py,sha256=vhU47Vcp_30M8tWu0FumbJ6pdUnFbBZm-ABdNlup__U,3821
|
14
|
-
nat/plugins/mcp/auth/register.py,sha256=yzphsn1I4a5G39_IacbuX0ZQqGM8fevvTUM_B94UXKE,1211
|
15
|
-
nvidia_nat_mcp-1.3.0rc1.dist-info/METADATA,sha256=NdfOVgW-bDqAZWInj84TQgDjH05tlBwEYO-mZdlwklM,1986
|
16
|
-
nvidia_nat_mcp-1.3.0rc1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
|
17
|
-
nvidia_nat_mcp-1.3.0rc1.dist-info/entry_points.txt,sha256=rYvUp4i-klBr3bVNh7zYOPXret704vTjvCk1qd7FooI,97
|
18
|
-
nvidia_nat_mcp-1.3.0rc1.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
|
19
|
-
nvidia_nat_mcp-1.3.0rc1.dist-info/RECORD,,
|
File without changes
|
{nvidia_nat_mcp-1.3.0rc1.dist-info → nvidia_nat_mcp-1.4.0a20251008.dist-info}/entry_points.txt
RENAMED
File without changes
|
File without changes
|