kiln-ai 0.21.0__py3-none-any.whl → 0.22.0__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.
Potentially problematic release.
This version of kiln-ai might be problematic. Click here for more details.
- kiln_ai/adapters/extractors/litellm_extractor.py +52 -32
- kiln_ai/adapters/extractors/test_litellm_extractor.py +169 -71
- kiln_ai/adapters/ml_embedding_model_list.py +330 -28
- kiln_ai/adapters/ml_model_list.py +503 -23
- kiln_ai/adapters/model_adapters/litellm_adapter.py +34 -7
- kiln_ai/adapters/model_adapters/test_litellm_adapter.py +78 -0
- kiln_ai/adapters/model_adapters/test_litellm_adapter_tools.py +119 -5
- kiln_ai/adapters/model_adapters/test_saving_adapter_results.py +9 -3
- kiln_ai/adapters/model_adapters/test_structured_output.py +6 -9
- kiln_ai/adapters/test_ml_embedding_model_list.py +89 -279
- kiln_ai/adapters/test_ml_model_list.py +0 -10
- kiln_ai/datamodel/basemodel.py +31 -3
- kiln_ai/datamodel/external_tool_server.py +206 -54
- kiln_ai/datamodel/extraction.py +14 -0
- kiln_ai/datamodel/task.py +5 -0
- kiln_ai/datamodel/task_output.py +41 -11
- kiln_ai/datamodel/test_attachment.py +3 -3
- kiln_ai/datamodel/test_basemodel.py +269 -13
- kiln_ai/datamodel/test_datasource.py +50 -0
- kiln_ai/datamodel/test_external_tool_server.py +534 -152
- kiln_ai/datamodel/test_extraction_model.py +31 -0
- kiln_ai/datamodel/test_task.py +35 -1
- kiln_ai/datamodel/test_tool_id.py +106 -1
- kiln_ai/datamodel/tool_id.py +36 -0
- kiln_ai/tools/base_tool.py +12 -3
- kiln_ai/tools/built_in_tools/math_tools.py +12 -4
- kiln_ai/tools/kiln_task_tool.py +158 -0
- kiln_ai/tools/mcp_server_tool.py +2 -2
- kiln_ai/tools/mcp_session_manager.py +50 -24
- kiln_ai/tools/rag_tools.py +12 -5
- kiln_ai/tools/test_kiln_task_tool.py +527 -0
- kiln_ai/tools/test_mcp_server_tool.py +4 -15
- kiln_ai/tools/test_mcp_session_manager.py +186 -226
- kiln_ai/tools/test_rag_tools.py +86 -5
- kiln_ai/tools/test_tool_registry.py +199 -5
- kiln_ai/tools/tool_registry.py +49 -17
- kiln_ai/utils/filesystem.py +4 -4
- kiln_ai/utils/open_ai_types.py +19 -2
- kiln_ai/utils/pdf_utils.py +21 -0
- kiln_ai/utils/test_open_ai_types.py +88 -12
- kiln_ai/utils/test_pdf_utils.py +14 -1
- {kiln_ai-0.21.0.dist-info → kiln_ai-0.22.0.dist-info}/METADATA +3 -1
- {kiln_ai-0.21.0.dist-info → kiln_ai-0.22.0.dist-info}/RECORD +45 -43
- {kiln_ai-0.21.0.dist-info → kiln_ai-0.22.0.dist-info}/WHEEL +0 -0
- {kiln_ai-0.21.0.dist-info → kiln_ai-0.22.0.dist-info}/licenses/LICENSE.txt +0 -0
|
@@ -8,11 +8,99 @@ from mcp.shared.exceptions import McpError
|
|
|
8
8
|
from mcp.types import ErrorData
|
|
9
9
|
from pydantic import ValidationError
|
|
10
10
|
|
|
11
|
-
from kiln_ai.datamodel.external_tool_server import
|
|
12
|
-
|
|
11
|
+
from kiln_ai.datamodel.external_tool_server import (
|
|
12
|
+
ExternalToolServer,
|
|
13
|
+
ToolServerType,
|
|
14
|
+
)
|
|
15
|
+
from kiln_ai.tools.mcp_session_manager import (
|
|
16
|
+
LOCAL_MCP_ERROR_INSTRUCTION,
|
|
17
|
+
MCPSessionManager,
|
|
18
|
+
)
|
|
13
19
|
from kiln_ai.utils.config import MCP_SECRETS_KEY
|
|
14
20
|
|
|
15
21
|
|
|
22
|
+
def create_remote_server(
|
|
23
|
+
headers=None,
|
|
24
|
+
secret_header_keys=None,
|
|
25
|
+
):
|
|
26
|
+
"""Factory function to create remote MCP servers with configurable properties."""
|
|
27
|
+
return ExternalToolServer(
|
|
28
|
+
name="test_server",
|
|
29
|
+
type=ToolServerType.remote_mcp,
|
|
30
|
+
description="Test server",
|
|
31
|
+
properties={
|
|
32
|
+
"server_url": "http://example.com/mcp",
|
|
33
|
+
"headers": headers or {},
|
|
34
|
+
"secret_header_keys": secret_header_keys or [],
|
|
35
|
+
},
|
|
36
|
+
)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def create_local_server(
|
|
40
|
+
command="python",
|
|
41
|
+
args=None,
|
|
42
|
+
env_vars=None,
|
|
43
|
+
secret_env_var_keys=None,
|
|
44
|
+
):
|
|
45
|
+
"""Factory function to create local MCP servers with configurable properties."""
|
|
46
|
+
return ExternalToolServer(
|
|
47
|
+
name="test_server",
|
|
48
|
+
type=ToolServerType.local_mcp,
|
|
49
|
+
description="Test server",
|
|
50
|
+
properties={
|
|
51
|
+
"command": command,
|
|
52
|
+
"args": args or [],
|
|
53
|
+
"env_vars": env_vars or {},
|
|
54
|
+
"secret_env_var_keys": secret_env_var_keys or [],
|
|
55
|
+
},
|
|
56
|
+
)
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
@pytest.fixture
|
|
60
|
+
def basic_remote_server():
|
|
61
|
+
return create_remote_server()
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
@pytest.fixture
|
|
65
|
+
def remote_server_with_auth():
|
|
66
|
+
return create_remote_server(headers={"Authorization": "Bearer token123"})
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
@pytest.fixture
|
|
70
|
+
def basic_local_server():
|
|
71
|
+
return create_local_server()
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
@pytest.fixture
|
|
75
|
+
def local_server_with_env():
|
|
76
|
+
"""Local MCP server with environment variables."""
|
|
77
|
+
return create_local_server(
|
|
78
|
+
args=["-m", "my_mcp_server"],
|
|
79
|
+
env_vars={"API_KEY": "test123"},
|
|
80
|
+
)
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
@pytest.fixture
|
|
84
|
+
def remote_server_with_secret_keys():
|
|
85
|
+
"""Remote MCP server with secret header keys."""
|
|
86
|
+
server = create_remote_server(
|
|
87
|
+
headers={"Content-Type": "application/json"},
|
|
88
|
+
secret_header_keys=["Authorization", "X-API-Key"],
|
|
89
|
+
)
|
|
90
|
+
return server
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
@pytest.fixture
|
|
94
|
+
def local_server_with_secret_keys():
|
|
95
|
+
"""Local MCP server with secret environment variable keys."""
|
|
96
|
+
server = create_local_server(
|
|
97
|
+
args=["-m", "my_server"],
|
|
98
|
+
env_vars={"PUBLIC_VAR": "public_value"},
|
|
99
|
+
secret_env_var_keys=["SECRET_API_KEY", "ANOTHER_SECRET"],
|
|
100
|
+
)
|
|
101
|
+
return server
|
|
102
|
+
|
|
103
|
+
|
|
16
104
|
class TestMCPSessionManager:
|
|
17
105
|
"""Unit tests for MCPSessionManager."""
|
|
18
106
|
|
|
@@ -84,8 +172,9 @@ class TestMCPSessionManager:
|
|
|
84
172
|
manager = MCPSessionManager()
|
|
85
173
|
|
|
86
174
|
# Create a mock exception-like object with exceptions attribute
|
|
87
|
-
class MockExceptionGroup:
|
|
175
|
+
class MockExceptionGroup(Exception):
|
|
88
176
|
def __init__(self, exceptions):
|
|
177
|
+
super().__init__("Mock exception group")
|
|
89
178
|
self.exceptions = exceptions
|
|
90
179
|
|
|
91
180
|
# Test finding target exception in exceptions list
|
|
@@ -105,8 +194,9 @@ class TestMCPSessionManager:
|
|
|
105
194
|
"""Test _extract_first_exception with nested objects having exceptions attribute."""
|
|
106
195
|
manager = MCPSessionManager()
|
|
107
196
|
|
|
108
|
-
class MockExceptionGroup:
|
|
197
|
+
class MockExceptionGroup(Exception):
|
|
109
198
|
def __init__(self, exceptions):
|
|
199
|
+
super().__init__("Mock exception group")
|
|
110
200
|
self.exceptions = exceptions
|
|
111
201
|
|
|
112
202
|
# Create nested structure
|
|
@@ -131,7 +221,7 @@ class TestMCPSessionManager:
|
|
|
131
221
|
manager = MCPSessionManager()
|
|
132
222
|
|
|
133
223
|
# Object without exceptions attribute should return None
|
|
134
|
-
class MockObject:
|
|
224
|
+
class MockObject(Exception):
|
|
135
225
|
pass
|
|
136
226
|
|
|
137
227
|
mock_obj = MockObject()
|
|
@@ -142,7 +232,7 @@ class TestMCPSessionManager:
|
|
|
142
232
|
"""Test _extract_first_exception with object that has None exceptions attribute."""
|
|
143
233
|
manager = MCPSessionManager()
|
|
144
234
|
|
|
145
|
-
class MockObject:
|
|
235
|
+
class MockObject(Exception):
|
|
146
236
|
exceptions = None
|
|
147
237
|
|
|
148
238
|
mock_obj = MockObject()
|
|
@@ -153,8 +243,9 @@ class TestMCPSessionManager:
|
|
|
153
243
|
"""Test _extract_first_exception with empty exceptions list."""
|
|
154
244
|
manager = MCPSessionManager()
|
|
155
245
|
|
|
156
|
-
class MockExceptionGroup:
|
|
246
|
+
class MockExceptionGroup(Exception):
|
|
157
247
|
def __init__(self):
|
|
248
|
+
super().__init__("Mock exception group")
|
|
158
249
|
self.exceptions = []
|
|
159
250
|
|
|
160
251
|
mock_group = MockExceptionGroup()
|
|
@@ -162,7 +253,9 @@ class TestMCPSessionManager:
|
|
|
162
253
|
assert result is None
|
|
163
254
|
|
|
164
255
|
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
165
|
-
async def test_successful_session_creation(
|
|
256
|
+
async def test_successful_session_creation(
|
|
257
|
+
self, mock_client, remote_server_with_auth
|
|
258
|
+
):
|
|
166
259
|
"""Test successful MCP session creation with mocked client."""
|
|
167
260
|
# Mock the streams
|
|
168
261
|
mock_read_stream = MagicMock()
|
|
@@ -175,17 +268,6 @@ class TestMCPSessionManager:
|
|
|
175
268
|
None,
|
|
176
269
|
)
|
|
177
270
|
|
|
178
|
-
# Create a valid tool server
|
|
179
|
-
tool_server = ExternalToolServer(
|
|
180
|
-
name="test_server",
|
|
181
|
-
type=ToolServerType.remote_mcp,
|
|
182
|
-
description="Test server",
|
|
183
|
-
properties={
|
|
184
|
-
"server_url": "http://example.com/mcp",
|
|
185
|
-
"headers": {"Authorization": "Bearer token123"},
|
|
186
|
-
},
|
|
187
|
-
)
|
|
188
|
-
|
|
189
271
|
manager = MCPSessionManager.shared()
|
|
190
272
|
|
|
191
273
|
with patch(
|
|
@@ -196,7 +278,7 @@ class TestMCPSessionManager:
|
|
|
196
278
|
mock_session_instance
|
|
197
279
|
)
|
|
198
280
|
|
|
199
|
-
async with manager.mcp_client(
|
|
281
|
+
async with manager.mcp_client(remote_server_with_auth) as session:
|
|
200
282
|
# Verify session is returned
|
|
201
283
|
assert session is mock_session_instance
|
|
202
284
|
|
|
@@ -209,7 +291,7 @@ class TestMCPSessionManager:
|
|
|
209
291
|
)
|
|
210
292
|
|
|
211
293
|
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
212
|
-
async def test_session_with_empty_headers(self, mock_client):
|
|
294
|
+
async def test_session_with_empty_headers(self, mock_client, basic_remote_server):
|
|
213
295
|
"""Test session creation when empty headers dict is provided."""
|
|
214
296
|
# Mock the streams
|
|
215
297
|
mock_read_stream = MagicMock()
|
|
@@ -222,16 +304,10 @@ class TestMCPSessionManager:
|
|
|
222
304
|
None,
|
|
223
305
|
)
|
|
224
306
|
|
|
225
|
-
#
|
|
226
|
-
tool_server =
|
|
227
|
-
|
|
228
|
-
|
|
229
|
-
description="Server with empty headers",
|
|
230
|
-
properties={
|
|
231
|
-
"server_url": "http://example.com/mcp",
|
|
232
|
-
"headers": {}, # Empty headers dict is required by pydantic
|
|
233
|
-
},
|
|
234
|
-
)
|
|
307
|
+
# Use basic server with empty headers
|
|
308
|
+
tool_server = basic_remote_server
|
|
309
|
+
tool_server.name = "empty_headers_server"
|
|
310
|
+
tool_server.description = "Server with empty headers"
|
|
235
311
|
|
|
236
312
|
manager = MCPSessionManager.shared()
|
|
237
313
|
|
|
@@ -262,7 +338,7 @@ class TestMCPSessionManager:
|
|
|
262
338
|
)
|
|
263
339
|
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
264
340
|
async def test_remote_mcp_http_status_errors(
|
|
265
|
-
self, mock_client, status_code, reason_phrase
|
|
341
|
+
self, mock_client, status_code, reason_phrase, basic_remote_server
|
|
266
342
|
):
|
|
267
343
|
"""Test remote MCP session handles various HTTP status errors with simplified message."""
|
|
268
344
|
# Create HTTP error with specific status code
|
|
@@ -276,13 +352,6 @@ class TestMCPSessionManager:
|
|
|
276
352
|
# Mock client to raise the HTTP error
|
|
277
353
|
mock_client.return_value.__aenter__.side_effect = http_error
|
|
278
354
|
|
|
279
|
-
tool_server = ExternalToolServer(
|
|
280
|
-
name="test_server",
|
|
281
|
-
type=ToolServerType.remote_mcp,
|
|
282
|
-
description="Test server",
|
|
283
|
-
properties={"server_url": "http://example.com/mcp", "headers": {}},
|
|
284
|
-
)
|
|
285
|
-
|
|
286
355
|
manager = MCPSessionManager.shared()
|
|
287
356
|
|
|
288
357
|
# All HTTP errors should now use the simplified message format
|
|
@@ -290,7 +359,7 @@ class TestMCPSessionManager:
|
|
|
290
359
|
with pytest.raises(
|
|
291
360
|
ValueError, match=expected_pattern.replace("(", r"\(").replace(")", r"\)")
|
|
292
361
|
):
|
|
293
|
-
async with manager.mcp_client(
|
|
362
|
+
async with manager.mcp_client(basic_remote_server):
|
|
294
363
|
pass
|
|
295
364
|
|
|
296
365
|
@pytest.mark.parametrize(
|
|
@@ -304,7 +373,7 @@ class TestMCPSessionManager:
|
|
|
304
373
|
)
|
|
305
374
|
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
306
375
|
async def test_remote_mcp_connection_errors(
|
|
307
|
-
self, mock_client, connection_error_type, error_message
|
|
376
|
+
self, mock_client, connection_error_type, error_message, basic_remote_server
|
|
308
377
|
):
|
|
309
378
|
"""Test remote MCP session handles various connection errors with simplified message."""
|
|
310
379
|
# Create connection error
|
|
@@ -318,22 +387,17 @@ class TestMCPSessionManager:
|
|
|
318
387
|
# Mock client to raise the connection error
|
|
319
388
|
mock_client.return_value.__aenter__.side_effect = connection_error
|
|
320
389
|
|
|
321
|
-
tool_server = ExternalToolServer(
|
|
322
|
-
name="test_server",
|
|
323
|
-
type=ToolServerType.remote_mcp,
|
|
324
|
-
description="Test server",
|
|
325
|
-
properties={"server_url": "http://example.com/mcp", "headers": {}},
|
|
326
|
-
)
|
|
327
|
-
|
|
328
390
|
manager = MCPSessionManager.shared()
|
|
329
391
|
|
|
330
392
|
# All connection errors should use the simplified message format
|
|
331
393
|
with pytest.raises(RuntimeError, match="Unable to connect to MCP server"):
|
|
332
|
-
async with manager.mcp_client(
|
|
394
|
+
async with manager.mcp_client(basic_remote_server):
|
|
333
395
|
pass
|
|
334
396
|
|
|
335
397
|
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
336
|
-
async def test_remote_mcp_http_error_in_nested_exceptions(
|
|
398
|
+
async def test_remote_mcp_http_error_in_nested_exceptions(
|
|
399
|
+
self, mock_client, basic_remote_server
|
|
400
|
+
):
|
|
337
401
|
"""Test remote MCP session extracts HTTP error from nested exceptions."""
|
|
338
402
|
# Create HTTP error nested in a mock exception group
|
|
339
403
|
response = MagicMock()
|
|
@@ -353,24 +417,19 @@ class TestMCPSessionManager:
|
|
|
353
417
|
# Mock client to raise the nested exception
|
|
354
418
|
mock_client.return_value.__aenter__.side_effect = group_error
|
|
355
419
|
|
|
356
|
-
tool_server = ExternalToolServer(
|
|
357
|
-
name="test_server",
|
|
358
|
-
type=ToolServerType.remote_mcp,
|
|
359
|
-
description="Test server",
|
|
360
|
-
properties={"server_url": "http://example.com/mcp", "headers": {}},
|
|
361
|
-
)
|
|
362
|
-
|
|
363
420
|
manager = MCPSessionManager.shared()
|
|
364
421
|
|
|
365
422
|
# Should extract the HTTP error from the nested structure
|
|
366
423
|
with pytest.raises(
|
|
367
424
|
ValueError, match=r"The MCP server rejected the request. Status 401"
|
|
368
425
|
):
|
|
369
|
-
async with manager.mcp_client(
|
|
426
|
+
async with manager.mcp_client(basic_remote_server):
|
|
370
427
|
pass
|
|
371
428
|
|
|
372
429
|
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
373
|
-
async def test_remote_mcp_connection_error_in_nested_exceptions(
|
|
430
|
+
async def test_remote_mcp_connection_error_in_nested_exceptions(
|
|
431
|
+
self, mock_client, basic_remote_server
|
|
432
|
+
):
|
|
374
433
|
"""Test remote MCP session extracts connection error from nested exceptions."""
|
|
375
434
|
# Create connection error nested in mock exception group
|
|
376
435
|
connection_error = ConnectionError("Connection timeout")
|
|
@@ -385,44 +444,34 @@ class TestMCPSessionManager:
|
|
|
385
444
|
# Mock client to raise the nested exception
|
|
386
445
|
mock_client.return_value.__aenter__.side_effect = group_error
|
|
387
446
|
|
|
388
|
-
tool_server = ExternalToolServer(
|
|
389
|
-
name="test_server",
|
|
390
|
-
type=ToolServerType.remote_mcp,
|
|
391
|
-
description="Test server",
|
|
392
|
-
properties={"server_url": "http://example.com/mcp", "headers": {}},
|
|
393
|
-
)
|
|
394
|
-
|
|
395
447
|
manager = MCPSessionManager.shared()
|
|
396
448
|
|
|
397
449
|
# Should extract the connection error from the nested structure
|
|
398
450
|
with pytest.raises(RuntimeError, match="Unable to connect to MCP server"):
|
|
399
|
-
async with manager.mcp_client(
|
|
451
|
+
async with manager.mcp_client(basic_remote_server):
|
|
400
452
|
pass
|
|
401
453
|
|
|
402
454
|
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
403
|
-
async def test_remote_mcp_unknown_error_fallback(
|
|
455
|
+
async def test_remote_mcp_unknown_error_fallback(
|
|
456
|
+
self, mock_client, basic_remote_server
|
|
457
|
+
):
|
|
404
458
|
"""Test remote MCP session handles unknown errors with fallback message."""
|
|
405
459
|
# Mock client to raise an unknown error type
|
|
406
460
|
unknown_error = RuntimeError("Unexpected error")
|
|
407
461
|
mock_client.return_value.__aenter__.side_effect = unknown_error
|
|
408
462
|
|
|
409
|
-
tool_server = ExternalToolServer(
|
|
410
|
-
name="test_server",
|
|
411
|
-
type=ToolServerType.remote_mcp,
|
|
412
|
-
description="Test server",
|
|
413
|
-
properties={"server_url": "http://example.com/mcp", "headers": {}},
|
|
414
|
-
)
|
|
415
|
-
|
|
416
463
|
manager = MCPSessionManager.shared()
|
|
417
464
|
|
|
418
465
|
# Should use the fallback error message
|
|
419
466
|
with pytest.raises(RuntimeError, match="Failed to connect to the MCP Server"):
|
|
420
|
-
async with manager.mcp_client(
|
|
467
|
+
async with manager.mcp_client(basic_remote_server):
|
|
421
468
|
pass
|
|
422
469
|
|
|
423
470
|
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
424
471
|
@patch("kiln_ai.utils.config.Config.shared")
|
|
425
|
-
async def test_session_with_secret_headers(
|
|
472
|
+
async def test_session_with_secret_headers(
|
|
473
|
+
self, mock_config, mock_client, remote_server_with_secret_keys
|
|
474
|
+
):
|
|
426
475
|
"""Test session creation with secret headers retrieved from config."""
|
|
427
476
|
# Mock the streams
|
|
428
477
|
mock_read_stream = MagicMock()
|
|
@@ -444,19 +493,10 @@ class TestMCPSessionManager:
|
|
|
444
493
|
}
|
|
445
494
|
mock_config.return_value = mock_config_instance
|
|
446
495
|
|
|
447
|
-
|
|
448
|
-
tool_server =
|
|
449
|
-
|
|
450
|
-
|
|
451
|
-
description="Server with secret headers",
|
|
452
|
-
properties={
|
|
453
|
-
"server_url": "http://example.com/mcp",
|
|
454
|
-
"headers": {"Content-Type": "application/json"},
|
|
455
|
-
"secret_header_keys": ["Authorization", "X-API-Key"],
|
|
456
|
-
},
|
|
457
|
-
)
|
|
458
|
-
# Set the server ID to match our mock secrets
|
|
459
|
-
tool_server.id = "test_server_id"
|
|
496
|
+
tool_server = remote_server_with_secret_keys
|
|
497
|
+
tool_server.id = "test_server_id" # Set the id
|
|
498
|
+
tool_server.name = "secret_headers_server"
|
|
499
|
+
tool_server.description = "Server with secret headers"
|
|
460
500
|
|
|
461
501
|
manager = MCPSessionManager.shared()
|
|
462
502
|
|
|
@@ -726,7 +766,7 @@ class TestMCPSessionManager:
|
|
|
726
766
|
tool_server.id = "test_server_id"
|
|
727
767
|
|
|
728
768
|
# Store original headers for comparison
|
|
729
|
-
original_headers = tool_server.properties
|
|
769
|
+
original_headers = tool_server.properties.get("headers", {}).copy()
|
|
730
770
|
|
|
731
771
|
manager = MCPSessionManager.shared()
|
|
732
772
|
|
|
@@ -743,18 +783,20 @@ class TestMCPSessionManager:
|
|
|
743
783
|
assert session is mock_session_instance
|
|
744
784
|
|
|
745
785
|
# Check that original headers are unchanged after first use
|
|
746
|
-
|
|
747
|
-
assert
|
|
748
|
-
assert "
|
|
786
|
+
headers = tool_server.properties.get("headers", {})
|
|
787
|
+
assert headers == original_headers
|
|
788
|
+
assert "Authorization" not in headers
|
|
789
|
+
assert "X-API-Key" not in headers
|
|
749
790
|
|
|
750
791
|
# Use the session a second time to ensure the bug doesn't occur on subsequent uses
|
|
751
792
|
async with manager.mcp_client(tool_server) as session:
|
|
752
793
|
assert session is mock_session_instance
|
|
753
794
|
|
|
754
795
|
# Check that original headers are still unchanged after second use
|
|
755
|
-
|
|
756
|
-
assert
|
|
757
|
-
assert "
|
|
796
|
+
headers = tool_server.properties.get("headers", {})
|
|
797
|
+
assert headers == original_headers
|
|
798
|
+
assert "Authorization" not in headers
|
|
799
|
+
assert "X-API-Key" not in headers
|
|
758
800
|
|
|
759
801
|
# Verify streamablehttp_client was called with merged headers both times
|
|
760
802
|
expected_headers = {
|
|
@@ -807,7 +849,7 @@ class TestMCPSessionManager:
|
|
|
807
849
|
tool_server.id = "test_server_id"
|
|
808
850
|
|
|
809
851
|
# Store original headers for comparison
|
|
810
|
-
original_headers = tool_server.properties
|
|
852
|
+
original_headers = tool_server.properties.get("headers", {}).copy()
|
|
811
853
|
|
|
812
854
|
# Simulate the buggy behavior by directly modifying the headers
|
|
813
855
|
# (This is what would happen without the .copy() fix)
|
|
@@ -825,18 +867,17 @@ class TestMCPSessionManager:
|
|
|
825
867
|
buggy_headers[header_name] = header_value
|
|
826
868
|
|
|
827
869
|
# Now the original properties would be contaminated with secrets!
|
|
828
|
-
|
|
829
|
-
assert
|
|
830
|
-
|
|
831
|
-
== "Bearer secret-token-123"
|
|
832
|
-
)
|
|
870
|
+
headers = tool_server.properties.get("headers", {})
|
|
871
|
+
assert "Authorization" in headers
|
|
872
|
+
assert headers["Authorization"] == "Bearer secret-token-123"
|
|
833
873
|
|
|
834
874
|
# This demonstrates the security bug - secrets are now permanently stored
|
|
835
875
|
# in the tool server properties and would be serialized/saved
|
|
836
|
-
|
|
876
|
+
headers = tool_server.properties.get("headers", {})
|
|
877
|
+
assert headers != original_headers
|
|
837
878
|
|
|
838
879
|
@patch("kiln_ai.tools.mcp_session_manager.stdio_client")
|
|
839
|
-
async def test_local_mcp_session_creation(self, mock_client):
|
|
880
|
+
async def test_local_mcp_session_creation(self, mock_client, local_server_with_env):
|
|
840
881
|
"""Test successful local MCP session creation with mocked client."""
|
|
841
882
|
# Mock the streams
|
|
842
883
|
mock_read_stream = MagicMock()
|
|
@@ -848,18 +889,6 @@ class TestMCPSessionManager:
|
|
|
848
889
|
mock_write_stream,
|
|
849
890
|
)
|
|
850
891
|
|
|
851
|
-
# Create a valid local tool server
|
|
852
|
-
tool_server = ExternalToolServer(
|
|
853
|
-
name="test_local_server",
|
|
854
|
-
type=ToolServerType.local_mcp,
|
|
855
|
-
description="Test local server",
|
|
856
|
-
properties={
|
|
857
|
-
"command": "python",
|
|
858
|
-
"args": ["-m", "my_mcp_server"],
|
|
859
|
-
"env_vars": {"API_KEY": "test123"},
|
|
860
|
-
},
|
|
861
|
-
)
|
|
862
|
-
|
|
863
892
|
manager = MCPSessionManager.shared()
|
|
864
893
|
|
|
865
894
|
with patch(
|
|
@@ -870,7 +899,7 @@ class TestMCPSessionManager:
|
|
|
870
899
|
mock_session_instance
|
|
871
900
|
)
|
|
872
901
|
|
|
873
|
-
async with manager.mcp_client(
|
|
902
|
+
async with manager.mcp_client(local_server_with_env) as session:
|
|
874
903
|
# Verify session is returned
|
|
875
904
|
assert session is mock_session_instance
|
|
876
905
|
|
|
@@ -937,36 +966,19 @@ class TestMCPSessionManager:
|
|
|
937
966
|
"""Test that missing command raises ValueError for local MCP."""
|
|
938
967
|
with pytest.raises(
|
|
939
968
|
ValidationError,
|
|
940
|
-
match="command must be a string
|
|
969
|
+
match="command must be a non-empty string",
|
|
941
970
|
):
|
|
942
971
|
ExternalToolServer(
|
|
943
972
|
name="missing_command_server",
|
|
944
973
|
type=ToolServerType.local_mcp,
|
|
945
974
|
description="Server missing command",
|
|
946
975
|
properties={
|
|
947
|
-
#
|
|
976
|
+
"command": "", # Empty command to trigger validation error
|
|
948
977
|
"args": ["arg1"],
|
|
949
978
|
"env_vars": {},
|
|
950
979
|
},
|
|
951
980
|
)
|
|
952
981
|
|
|
953
|
-
async def test_local_mcp_missing_args_error(self):
|
|
954
|
-
"""Test that missing args raises ValueError for local MCP."""
|
|
955
|
-
with pytest.raises(
|
|
956
|
-
ValidationError,
|
|
957
|
-
match="arguments must be a list to start a local MCP server",
|
|
958
|
-
):
|
|
959
|
-
ExternalToolServer(
|
|
960
|
-
name="missing_args_server",
|
|
961
|
-
type=ToolServerType.local_mcp,
|
|
962
|
-
description="Server missing args",
|
|
963
|
-
properties={
|
|
964
|
-
"command": "python",
|
|
965
|
-
# No args provided
|
|
966
|
-
"env_vars": {},
|
|
967
|
-
},
|
|
968
|
-
)
|
|
969
|
-
|
|
970
982
|
async def test_local_mcp_empty_args_allowed(self):
|
|
971
983
|
"""Test that empty args list is now allowed for local MCP."""
|
|
972
984
|
# Should not raise any exception - empty args are now allowed
|
|
@@ -983,63 +995,14 @@ class TestMCPSessionManager:
|
|
|
983
995
|
|
|
984
996
|
assert tool_server.name == "empty_args_server"
|
|
985
997
|
assert tool_server.type == ToolServerType.local_mcp
|
|
986
|
-
|
|
987
|
-
|
|
988
|
-
async def test_local_mcp_session_empty_command_runtime_error(self):
|
|
989
|
-
"""Test that empty command string raises ValueError during session creation."""
|
|
990
|
-
# Create a valid tool server first
|
|
991
|
-
tool_server = ExternalToolServer(
|
|
992
|
-
name="test_server",
|
|
993
|
-
type=ToolServerType.local_mcp,
|
|
994
|
-
description="Test server",
|
|
995
|
-
properties={
|
|
996
|
-
"command": "python",
|
|
997
|
-
"args": ["arg1"],
|
|
998
|
-
"env_vars": {},
|
|
999
|
-
},
|
|
1000
|
-
)
|
|
1001
|
-
|
|
1002
|
-
# Manually modify the properties after creation to bypass pydantic validation
|
|
1003
|
-
tool_server.properties["command"] = ""
|
|
1004
|
-
|
|
1005
|
-
manager = MCPSessionManager.shared()
|
|
1006
|
-
|
|
1007
|
-
with pytest.raises(
|
|
1008
|
-
ValueError,
|
|
1009
|
-
match="Attempted to start local MCP server, but no command was provided",
|
|
1010
|
-
):
|
|
1011
|
-
async with manager.mcp_client(tool_server):
|
|
1012
|
-
pass
|
|
1013
|
-
|
|
1014
|
-
async def test_local_mcp_session_invalid_args_type_runtime_error(self):
|
|
1015
|
-
"""Test that non-list args raises ValueError during session creation."""
|
|
1016
|
-
# Create a valid tool server first
|
|
1017
|
-
tool_server = ExternalToolServer(
|
|
1018
|
-
name="test_server",
|
|
1019
|
-
type=ToolServerType.local_mcp,
|
|
1020
|
-
description="Test server",
|
|
1021
|
-
properties={
|
|
1022
|
-
"command": "python",
|
|
1023
|
-
"args": ["arg1"],
|
|
1024
|
-
"env_vars": {},
|
|
1025
|
-
},
|
|
1026
|
-
)
|
|
1027
|
-
|
|
1028
|
-
# Manually modify the properties after creation to bypass pydantic validation
|
|
1029
|
-
tool_server.properties["args"] = "not a list"
|
|
1030
|
-
|
|
1031
|
-
manager = MCPSessionManager.shared()
|
|
1032
|
-
|
|
1033
|
-
with pytest.raises(
|
|
1034
|
-
ValueError,
|
|
1035
|
-
match="Attempted to start local MCP server, but args is not a list of strings",
|
|
1036
|
-
):
|
|
1037
|
-
async with manager.mcp_client(tool_server):
|
|
1038
|
-
pass
|
|
998
|
+
args = tool_server.properties.get("args", [])
|
|
999
|
+
assert args == []
|
|
1039
1000
|
|
|
1040
1001
|
@patch("kiln_ai.tools.mcp_session_manager.stdio_client")
|
|
1041
1002
|
@patch("kiln_ai.utils.config.Config.shared")
|
|
1042
|
-
async def test_local_mcp_session_with_secrets(
|
|
1003
|
+
async def test_local_mcp_session_with_secrets(
|
|
1004
|
+
self, mock_config, mock_client, local_server_with_secret_keys
|
|
1005
|
+
):
|
|
1043
1006
|
"""Test local MCP session creation with secret environment variables."""
|
|
1044
1007
|
# Mock config to return different values based on the key
|
|
1045
1008
|
mock_config_instance = MagicMock()
|
|
@@ -1067,19 +1030,7 @@ class TestMCPSessionManager:
|
|
|
1067
1030
|
mock_write_stream,
|
|
1068
1031
|
)
|
|
1069
1032
|
|
|
1070
|
-
|
|
1071
|
-
tool_server = ExternalToolServer(
|
|
1072
|
-
name="test_server",
|
|
1073
|
-
type=ToolServerType.local_mcp,
|
|
1074
|
-
description="Test server with secrets",
|
|
1075
|
-
properties={
|
|
1076
|
-
"command": "python",
|
|
1077
|
-
"args": ["-m", "my_server"],
|
|
1078
|
-
"env_vars": {"PUBLIC_VAR": "public_value"},
|
|
1079
|
-
"secret_env_var_keys": ["SECRET_API_KEY", "ANOTHER_SECRET"],
|
|
1080
|
-
},
|
|
1081
|
-
)
|
|
1082
|
-
# Set the server ID to match our mock secrets
|
|
1033
|
+
tool_server = local_server_with_secret_keys
|
|
1083
1034
|
tool_server.id = "test_server_id"
|
|
1084
1035
|
|
|
1085
1036
|
manager = MCPSessionManager.shared()
|
|
@@ -1139,7 +1090,7 @@ class TestMCPSessionManager:
|
|
|
1139
1090
|
)
|
|
1140
1091
|
@patch("kiln_ai.tools.mcp_session_manager.stdio_client")
|
|
1141
1092
|
async def test_local_mcp_various_errors_use_simplified_message(
|
|
1142
|
-
self, mock_client, error_type, error_message
|
|
1093
|
+
self, mock_client, error_type, error_message, basic_local_server
|
|
1143
1094
|
):
|
|
1144
1095
|
"""Test local MCP session handles various errors with simplified message."""
|
|
1145
1096
|
# Create the appropriate error
|
|
@@ -1152,22 +1103,11 @@ class TestMCPSessionManager:
|
|
|
1152
1103
|
# Mock client to raise the error
|
|
1153
1104
|
mock_client.return_value.__aenter__.side_effect = test_error
|
|
1154
1105
|
|
|
1155
|
-
tool_server = ExternalToolServer(
|
|
1156
|
-
name="test_server",
|
|
1157
|
-
type=ToolServerType.local_mcp,
|
|
1158
|
-
description="Test server",
|
|
1159
|
-
properties={
|
|
1160
|
-
"command": "python",
|
|
1161
|
-
"args": ["-m", "my_server"],
|
|
1162
|
-
"env_vars": {},
|
|
1163
|
-
},
|
|
1164
|
-
)
|
|
1165
|
-
|
|
1166
1106
|
manager = MCPSessionManager.shared()
|
|
1167
1107
|
|
|
1168
1108
|
# All local errors should now use the simplified message format
|
|
1169
|
-
with pytest.raises(RuntimeError, match=
|
|
1170
|
-
async with manager.mcp_client(
|
|
1109
|
+
with pytest.raises(RuntimeError, match=LOCAL_MCP_ERROR_INSTRUCTION):
|
|
1110
|
+
async with manager.mcp_client(basic_local_server):
|
|
1171
1111
|
pass
|
|
1172
1112
|
|
|
1173
1113
|
@patch("kiln_ai.tools.mcp_session_manager.stdio_client")
|
|
@@ -1201,7 +1141,7 @@ class TestMCPSessionManager:
|
|
|
1201
1141
|
manager = MCPSessionManager.shared()
|
|
1202
1142
|
|
|
1203
1143
|
# Should extract the McpError from the nested structure and use simplified message
|
|
1204
|
-
with pytest.raises(RuntimeError, match=
|
|
1144
|
+
with pytest.raises(RuntimeError, match=LOCAL_MCP_ERROR_INSTRUCTION):
|
|
1205
1145
|
async with manager.mcp_client(tool_server):
|
|
1206
1146
|
pass
|
|
1207
1147
|
|
|
@@ -1219,14 +1159,10 @@ class TestMCPSessionManager:
|
|
|
1219
1159
|
|
|
1220
1160
|
for original_error in test_exceptions:
|
|
1221
1161
|
with pytest.raises(RuntimeError) as exc_info:
|
|
1222
|
-
manager._raise_local_mcp_error(original_error)
|
|
1162
|
+
manager._raise_local_mcp_error(original_error, "")
|
|
1223
1163
|
|
|
1224
1164
|
# Check that the error message contains expected text
|
|
1225
|
-
assert
|
|
1226
|
-
assert (
|
|
1227
|
-
"Please verify your command, arguments, and environment variables"
|
|
1228
|
-
in str(exc_info.value)
|
|
1229
|
-
)
|
|
1165
|
+
assert LOCAL_MCP_ERROR_INSTRUCTION in str(exc_info.value)
|
|
1230
1166
|
assert str(original_error) in str(exc_info.value)
|
|
1231
1167
|
|
|
1232
1168
|
# Check that the original exception is chained
|
|
@@ -1531,6 +1467,31 @@ class TestMCPSessionManager:
|
|
|
1531
1467
|
assert result2 == "/cached/path"
|
|
1532
1468
|
assert mock_run.call_count == 1 # Should not have been called again
|
|
1533
1469
|
|
|
1470
|
+
async def test_mcp_client_with_kiln_task_raises_error(self):
|
|
1471
|
+
"""Test that mcp_client raises ValueError when passed a Kiln task tool server."""
|
|
1472
|
+
# Create a Kiln task tool server
|
|
1473
|
+
kiln_task_server = ExternalToolServer(
|
|
1474
|
+
name="test_kiln_task",
|
|
1475
|
+
type=ToolServerType.kiln_task,
|
|
1476
|
+
description="Test Kiln task",
|
|
1477
|
+
properties={
|
|
1478
|
+
"task_id": "task_123",
|
|
1479
|
+
"run_config_id": "config_456",
|
|
1480
|
+
"name": "test_task",
|
|
1481
|
+
"description": "A test task for validation",
|
|
1482
|
+
"is_archived": False,
|
|
1483
|
+
},
|
|
1484
|
+
)
|
|
1485
|
+
|
|
1486
|
+
manager = MCPSessionManager.shared()
|
|
1487
|
+
|
|
1488
|
+
# Should raise ValueError with specific message
|
|
1489
|
+
with pytest.raises(
|
|
1490
|
+
ValueError, match="Kiln task tools are not available from an MCP server"
|
|
1491
|
+
):
|
|
1492
|
+
async with manager.mcp_client(kiln_task_server):
|
|
1493
|
+
pass
|
|
1494
|
+
|
|
1534
1495
|
|
|
1535
1496
|
class TestMCPServerIntegration:
|
|
1536
1497
|
"""Integration tests for MCPServer using real services."""
|
|
@@ -1546,7 +1507,6 @@ class TestMCPServerIntegration:
|
|
|
1546
1507
|
description="Postman Echo MCP Server for testing",
|
|
1547
1508
|
properties={
|
|
1548
1509
|
"server_url": "https://postman-echo-mcp.fly.dev/",
|
|
1549
|
-
"headers": {},
|
|
1550
1510
|
},
|
|
1551
1511
|
)
|
|
1552
1512
|
|