kiln-ai 0.19.0__py3-none-any.whl → 0.20.1__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/__init__.py +2 -2
- kiln_ai/adapters/adapter_registry.py +19 -1
- kiln_ai/adapters/chat/chat_formatter.py +8 -12
- kiln_ai/adapters/chat/test_chat_formatter.py +6 -2
- kiln_ai/adapters/docker_model_runner_tools.py +119 -0
- kiln_ai/adapters/eval/base_eval.py +2 -2
- kiln_ai/adapters/eval/eval_runner.py +3 -1
- kiln_ai/adapters/eval/g_eval.py +2 -2
- kiln_ai/adapters/eval/test_base_eval.py +1 -1
- kiln_ai/adapters/eval/test_g_eval.py +3 -4
- kiln_ai/adapters/fine_tune/__init__.py +1 -1
- kiln_ai/adapters/fine_tune/openai_finetune.py +14 -4
- kiln_ai/adapters/fine_tune/test_openai_finetune.py +108 -111
- kiln_ai/adapters/ml_model_list.py +380 -34
- kiln_ai/adapters/model_adapters/base_adapter.py +51 -21
- kiln_ai/adapters/model_adapters/litellm_adapter.py +383 -79
- kiln_ai/adapters/model_adapters/test_base_adapter.py +193 -17
- kiln_ai/adapters/model_adapters/test_litellm_adapter.py +406 -1
- kiln_ai/adapters/model_adapters/test_litellm_adapter_tools.py +1103 -0
- kiln_ai/adapters/model_adapters/test_saving_adapter_results.py +5 -5
- kiln_ai/adapters/model_adapters/test_structured_output.py +110 -4
- kiln_ai/adapters/parsers/__init__.py +1 -1
- kiln_ai/adapters/provider_tools.py +15 -1
- kiln_ai/adapters/repair/test_repair_task.py +12 -9
- kiln_ai/adapters/run_output.py +3 -0
- kiln_ai/adapters/test_adapter_registry.py +80 -1
- kiln_ai/adapters/test_docker_model_runner_tools.py +305 -0
- kiln_ai/adapters/test_ml_model_list.py +39 -1
- kiln_ai/adapters/test_prompt_adaptors.py +13 -6
- kiln_ai/adapters/test_provider_tools.py +55 -0
- kiln_ai/adapters/test_remote_config.py +98 -0
- kiln_ai/datamodel/__init__.py +23 -21
- kiln_ai/datamodel/datamodel_enums.py +1 -0
- kiln_ai/datamodel/eval.py +1 -1
- kiln_ai/datamodel/external_tool_server.py +298 -0
- kiln_ai/datamodel/json_schema.py +25 -10
- kiln_ai/datamodel/project.py +8 -1
- kiln_ai/datamodel/registry.py +0 -15
- kiln_ai/datamodel/run_config.py +62 -0
- kiln_ai/datamodel/task.py +2 -77
- kiln_ai/datamodel/task_output.py +6 -1
- kiln_ai/datamodel/task_run.py +41 -0
- kiln_ai/datamodel/test_basemodel.py +3 -3
- kiln_ai/datamodel/test_example_models.py +175 -0
- kiln_ai/datamodel/test_external_tool_server.py +691 -0
- kiln_ai/datamodel/test_registry.py +8 -3
- kiln_ai/datamodel/test_task.py +15 -47
- kiln_ai/datamodel/test_tool_id.py +239 -0
- kiln_ai/datamodel/tool_id.py +83 -0
- kiln_ai/tools/__init__.py +8 -0
- kiln_ai/tools/base_tool.py +82 -0
- kiln_ai/tools/built_in_tools/__init__.py +13 -0
- kiln_ai/tools/built_in_tools/math_tools.py +124 -0
- kiln_ai/tools/built_in_tools/test_math_tools.py +204 -0
- kiln_ai/tools/mcp_server_tool.py +95 -0
- kiln_ai/tools/mcp_session_manager.py +243 -0
- kiln_ai/tools/test_base_tools.py +199 -0
- kiln_ai/tools/test_mcp_server_tool.py +457 -0
- kiln_ai/tools/test_mcp_session_manager.py +1585 -0
- kiln_ai/tools/test_tool_registry.py +473 -0
- kiln_ai/tools/tool_registry.py +64 -0
- kiln_ai/utils/config.py +22 -0
- kiln_ai/utils/open_ai_types.py +94 -0
- kiln_ai/utils/project_utils.py +17 -0
- kiln_ai/utils/test_config.py +138 -1
- kiln_ai/utils/test_open_ai_types.py +131 -0
- {kiln_ai-0.19.0.dist-info → kiln_ai-0.20.1.dist-info}/METADATA +6 -5
- {kiln_ai-0.19.0.dist-info → kiln_ai-0.20.1.dist-info}/RECORD +70 -47
- {kiln_ai-0.19.0.dist-info → kiln_ai-0.20.1.dist-info}/WHEEL +0 -0
- {kiln_ai-0.19.0.dist-info → kiln_ai-0.20.1.dist-info}/licenses/LICENSE.txt +0 -0
|
@@ -0,0 +1,1585 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import subprocess
|
|
3
|
+
from unittest.mock import AsyncMock, MagicMock, patch
|
|
4
|
+
|
|
5
|
+
import httpx
|
|
6
|
+
import pytest
|
|
7
|
+
from mcp.shared.exceptions import McpError
|
|
8
|
+
from mcp.types import ErrorData
|
|
9
|
+
from pydantic import ValidationError
|
|
10
|
+
|
|
11
|
+
from kiln_ai.datamodel.external_tool_server import ExternalToolServer, ToolServerType
|
|
12
|
+
from kiln_ai.tools.mcp_session_manager import MCPSessionManager
|
|
13
|
+
from kiln_ai.utils.config import MCP_SECRETS_KEY
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class TestMCPSessionManager:
|
|
17
|
+
"""Unit tests for MCPSessionManager."""
|
|
18
|
+
|
|
19
|
+
def test_singleton_behavior(self):
|
|
20
|
+
"""Test that MCPSessionManager follows singleton pattern."""
|
|
21
|
+
# Get two instances
|
|
22
|
+
instance1 = MCPSessionManager.shared()
|
|
23
|
+
instance2 = MCPSessionManager.shared()
|
|
24
|
+
|
|
25
|
+
# They should be the same object
|
|
26
|
+
assert instance1 is instance2
|
|
27
|
+
assert id(instance1) == id(instance2)
|
|
28
|
+
|
|
29
|
+
def test_singleton_reset_for_testing(self):
|
|
30
|
+
"""Test that we can reset the singleton for testing purposes."""
|
|
31
|
+
# Get an instance
|
|
32
|
+
instance1 = MCPSessionManager.shared()
|
|
33
|
+
|
|
34
|
+
# Reset the singleton
|
|
35
|
+
MCPSessionManager._shared_instance = None
|
|
36
|
+
|
|
37
|
+
# Get a new instance
|
|
38
|
+
instance2 = MCPSessionManager.shared()
|
|
39
|
+
|
|
40
|
+
# They should be different objects
|
|
41
|
+
assert instance1 is not instance2
|
|
42
|
+
|
|
43
|
+
# Note: Testing invalid tool server types is not possible because:
|
|
44
|
+
# 1. The ToolServerType enum only has one value: remote_mcp
|
|
45
|
+
# 2. Pydantic validation prevents creating objects with invalid types
|
|
46
|
+
# 3. Pydantic prevents modifying the type field to invalid values after creation
|
|
47
|
+
# The RuntimeError check in MCPSessionManager.mcp_client is defensive programming
|
|
48
|
+
# that would only be triggered if new enum values are added without updating the match statement.
|
|
49
|
+
|
|
50
|
+
@pytest.mark.parametrize(
|
|
51
|
+
"exception,target_type,expected_result",
|
|
52
|
+
[
|
|
53
|
+
# Direct matches
|
|
54
|
+
(ValueError("test"), ValueError, True),
|
|
55
|
+
(ConnectionError("conn"), ConnectionError, True),
|
|
56
|
+
(FileNotFoundError("file"), FileNotFoundError, True),
|
|
57
|
+
# Non-matches
|
|
58
|
+
(ValueError("test"), TypeError, False),
|
|
59
|
+
(ConnectionError("conn"), ValueError, False),
|
|
60
|
+
# Tuple targets - matches
|
|
61
|
+
(ValueError("test"), (ValueError, TypeError), True),
|
|
62
|
+
(ConnectionError("conn"), (ValueError, ConnectionError), True),
|
|
63
|
+
# Tuple targets - non-matches
|
|
64
|
+
(RuntimeError("test"), (ValueError, TypeError), False),
|
|
65
|
+
# Inheritance - FileNotFoundError is subclass of OSError
|
|
66
|
+
(FileNotFoundError("file"), OSError, True),
|
|
67
|
+
(ConnectionError("conn"), OSError, True),
|
|
68
|
+
],
|
|
69
|
+
)
|
|
70
|
+
def test_extract_first_exception_direct_cases(
|
|
71
|
+
self, exception, target_type, expected_result
|
|
72
|
+
):
|
|
73
|
+
"""Test _extract_first_exception with direct exception cases."""
|
|
74
|
+
manager = MCPSessionManager()
|
|
75
|
+
result = manager._extract_first_exception(exception, target_type)
|
|
76
|
+
|
|
77
|
+
if expected_result:
|
|
78
|
+
assert result is exception
|
|
79
|
+
else:
|
|
80
|
+
assert result is None
|
|
81
|
+
|
|
82
|
+
def test_extract_first_exception_with_exceptions_attribute(self):
|
|
83
|
+
"""Test _extract_first_exception with object that has exceptions attribute."""
|
|
84
|
+
manager = MCPSessionManager()
|
|
85
|
+
|
|
86
|
+
# Create a mock exception-like object with exceptions attribute
|
|
87
|
+
class MockExceptionGroup:
|
|
88
|
+
def __init__(self, exceptions):
|
|
89
|
+
self.exceptions = exceptions
|
|
90
|
+
|
|
91
|
+
# Test finding target exception in exceptions list
|
|
92
|
+
target_exception = ValueError("found")
|
|
93
|
+
mock_group = MockExceptionGroup(
|
|
94
|
+
[TypeError("other"), target_exception, RuntimeError("another")]
|
|
95
|
+
)
|
|
96
|
+
|
|
97
|
+
result = manager._extract_first_exception(mock_group, ValueError)
|
|
98
|
+
assert result is target_exception
|
|
99
|
+
|
|
100
|
+
# Test not finding target exception
|
|
101
|
+
result = manager._extract_first_exception(mock_group, KeyError)
|
|
102
|
+
assert result is None
|
|
103
|
+
|
|
104
|
+
def test_extract_first_exception_nested_exceptions_attribute(self):
|
|
105
|
+
"""Test _extract_first_exception with nested objects having exceptions attribute."""
|
|
106
|
+
manager = MCPSessionManager()
|
|
107
|
+
|
|
108
|
+
class MockExceptionGroup:
|
|
109
|
+
def __init__(self, exceptions):
|
|
110
|
+
self.exceptions = exceptions
|
|
111
|
+
|
|
112
|
+
# Create nested structure
|
|
113
|
+
target_exception = FileNotFoundError("nested target")
|
|
114
|
+
inner_group = MockExceptionGroup([target_exception, ValueError("inner other")])
|
|
115
|
+
outer_group = MockExceptionGroup([TypeError("outer other"), inner_group])
|
|
116
|
+
|
|
117
|
+
# Should find deeply nested exception
|
|
118
|
+
result = manager._extract_first_exception(outer_group, FileNotFoundError)
|
|
119
|
+
assert result is target_exception
|
|
120
|
+
|
|
121
|
+
# Should find by parent class
|
|
122
|
+
result = manager._extract_first_exception(outer_group, OSError)
|
|
123
|
+
assert result is target_exception
|
|
124
|
+
|
|
125
|
+
# Should not find non-existent exception
|
|
126
|
+
result = manager._extract_first_exception(outer_group, KeyError)
|
|
127
|
+
assert result is None
|
|
128
|
+
|
|
129
|
+
def test_extract_first_exception_no_exceptions_attribute(self):
|
|
130
|
+
"""Test _extract_first_exception with object that has no exceptions attribute."""
|
|
131
|
+
manager = MCPSessionManager()
|
|
132
|
+
|
|
133
|
+
# Object without exceptions attribute should return None
|
|
134
|
+
class MockObject:
|
|
135
|
+
pass
|
|
136
|
+
|
|
137
|
+
mock_obj = MockObject()
|
|
138
|
+
result = manager._extract_first_exception(mock_obj, ValueError)
|
|
139
|
+
assert result is None
|
|
140
|
+
|
|
141
|
+
def test_extract_first_exception_none_exceptions_attribute(self):
|
|
142
|
+
"""Test _extract_first_exception with object that has None exceptions attribute."""
|
|
143
|
+
manager = MCPSessionManager()
|
|
144
|
+
|
|
145
|
+
class MockObject:
|
|
146
|
+
exceptions = None
|
|
147
|
+
|
|
148
|
+
mock_obj = MockObject()
|
|
149
|
+
result = manager._extract_first_exception(mock_obj, ValueError)
|
|
150
|
+
assert result is None
|
|
151
|
+
|
|
152
|
+
def test_extract_first_exception_empty_exceptions_list(self):
|
|
153
|
+
"""Test _extract_first_exception with empty exceptions list."""
|
|
154
|
+
manager = MCPSessionManager()
|
|
155
|
+
|
|
156
|
+
class MockExceptionGroup:
|
|
157
|
+
def __init__(self):
|
|
158
|
+
self.exceptions = []
|
|
159
|
+
|
|
160
|
+
mock_group = MockExceptionGroup()
|
|
161
|
+
result = manager._extract_first_exception(mock_group, ValueError)
|
|
162
|
+
assert result is None
|
|
163
|
+
|
|
164
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
165
|
+
async def test_successful_session_creation(self, mock_client):
|
|
166
|
+
"""Test successful MCP session creation with mocked client."""
|
|
167
|
+
# Mock the streams
|
|
168
|
+
mock_read_stream = MagicMock()
|
|
169
|
+
mock_write_stream = MagicMock()
|
|
170
|
+
|
|
171
|
+
# Configure the mock client context manager
|
|
172
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
173
|
+
mock_read_stream,
|
|
174
|
+
mock_write_stream,
|
|
175
|
+
None,
|
|
176
|
+
)
|
|
177
|
+
|
|
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
|
+
manager = MCPSessionManager.shared()
|
|
190
|
+
|
|
191
|
+
with patch(
|
|
192
|
+
"kiln_ai.tools.mcp_session_manager.ClientSession"
|
|
193
|
+
) as mock_session_class:
|
|
194
|
+
mock_session_instance = AsyncMock()
|
|
195
|
+
mock_session_class.return_value.__aenter__.return_value = (
|
|
196
|
+
mock_session_instance
|
|
197
|
+
)
|
|
198
|
+
|
|
199
|
+
async with manager.mcp_client(tool_server) as session:
|
|
200
|
+
# Verify session is returned
|
|
201
|
+
assert session is mock_session_instance
|
|
202
|
+
|
|
203
|
+
# Verify initialize was called
|
|
204
|
+
mock_session_instance.initialize.assert_called_once()
|
|
205
|
+
|
|
206
|
+
# Verify streamablehttp_client was called with correct parameters
|
|
207
|
+
mock_client.assert_called_once_with(
|
|
208
|
+
"http://example.com/mcp", headers={"Authorization": "Bearer token123"}
|
|
209
|
+
)
|
|
210
|
+
|
|
211
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
212
|
+
async def test_session_with_empty_headers(self, mock_client):
|
|
213
|
+
"""Test session creation when empty headers dict is provided."""
|
|
214
|
+
# Mock the streams
|
|
215
|
+
mock_read_stream = MagicMock()
|
|
216
|
+
mock_write_stream = MagicMock()
|
|
217
|
+
|
|
218
|
+
# Configure the mock client context manager
|
|
219
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
220
|
+
mock_read_stream,
|
|
221
|
+
mock_write_stream,
|
|
222
|
+
None,
|
|
223
|
+
)
|
|
224
|
+
|
|
225
|
+
# Create a tool server with empty headers
|
|
226
|
+
tool_server = ExternalToolServer(
|
|
227
|
+
name="empty_headers_server",
|
|
228
|
+
type=ToolServerType.remote_mcp,
|
|
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
|
+
)
|
|
235
|
+
|
|
236
|
+
manager = MCPSessionManager.shared()
|
|
237
|
+
|
|
238
|
+
with patch(
|
|
239
|
+
"kiln_ai.tools.mcp_session_manager.ClientSession"
|
|
240
|
+
) as mock_session_class:
|
|
241
|
+
mock_session_instance = AsyncMock()
|
|
242
|
+
mock_session_class.return_value.__aenter__.return_value = (
|
|
243
|
+
mock_session_instance
|
|
244
|
+
)
|
|
245
|
+
|
|
246
|
+
async with manager.mcp_client(tool_server) as session:
|
|
247
|
+
assert session is mock_session_instance
|
|
248
|
+
|
|
249
|
+
# Verify streamablehttp_client was called with empty headers dict
|
|
250
|
+
mock_client.assert_called_once_with("http://example.com/mcp", headers={})
|
|
251
|
+
|
|
252
|
+
@pytest.mark.parametrize(
|
|
253
|
+
"status_code,reason_phrase",
|
|
254
|
+
[
|
|
255
|
+
(400, "Bad Request"),
|
|
256
|
+
(401, "Unauthorized"),
|
|
257
|
+
(403, "Forbidden"),
|
|
258
|
+
(404, "Not Found"),
|
|
259
|
+
(500, "Internal Server Error"),
|
|
260
|
+
(502, "Bad Gateway"),
|
|
261
|
+
],
|
|
262
|
+
)
|
|
263
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
264
|
+
async def test_remote_mcp_http_status_errors(
|
|
265
|
+
self, mock_client, status_code, reason_phrase
|
|
266
|
+
):
|
|
267
|
+
"""Test remote MCP session handles various HTTP status errors with simplified message."""
|
|
268
|
+
# Create HTTP error with specific status code
|
|
269
|
+
response = MagicMock()
|
|
270
|
+
response.status_code = status_code
|
|
271
|
+
response.reason_phrase = reason_phrase
|
|
272
|
+
http_error = httpx.HTTPStatusError(
|
|
273
|
+
reason_phrase, request=MagicMock(), response=response
|
|
274
|
+
)
|
|
275
|
+
|
|
276
|
+
# Mock client to raise the HTTP error
|
|
277
|
+
mock_client.return_value.__aenter__.side_effect = http_error
|
|
278
|
+
|
|
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
|
+
manager = MCPSessionManager.shared()
|
|
287
|
+
|
|
288
|
+
# All HTTP errors should now use the simplified message format
|
|
289
|
+
expected_pattern = f"The MCP server rejected the request. Status {status_code}. Response from server:\n{reason_phrase}"
|
|
290
|
+
with pytest.raises(
|
|
291
|
+
ValueError, match=expected_pattern.replace("(", r"\(").replace(")", r"\)")
|
|
292
|
+
):
|
|
293
|
+
async with manager.mcp_client(tool_server):
|
|
294
|
+
pass
|
|
295
|
+
|
|
296
|
+
@pytest.mark.parametrize(
|
|
297
|
+
"connection_error_type,error_message",
|
|
298
|
+
[
|
|
299
|
+
(ConnectionError, "Connection refused"),
|
|
300
|
+
(OSError, "Network is unreachable"),
|
|
301
|
+
(httpx.RequestError, "Request failed"),
|
|
302
|
+
(httpx.ConnectError, "Connection error"),
|
|
303
|
+
],
|
|
304
|
+
)
|
|
305
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
306
|
+
async def test_remote_mcp_connection_errors(
|
|
307
|
+
self, mock_client, connection_error_type, error_message
|
|
308
|
+
):
|
|
309
|
+
"""Test remote MCP session handles various connection errors with simplified message."""
|
|
310
|
+
# Create connection error
|
|
311
|
+
if connection_error_type == httpx.RequestError:
|
|
312
|
+
connection_error = connection_error_type(error_message, request=MagicMock())
|
|
313
|
+
elif connection_error_type == httpx.ConnectError:
|
|
314
|
+
connection_error = connection_error_type(error_message, request=MagicMock())
|
|
315
|
+
else:
|
|
316
|
+
connection_error = connection_error_type(error_message)
|
|
317
|
+
|
|
318
|
+
# Mock client to raise the connection error
|
|
319
|
+
mock_client.return_value.__aenter__.side_effect = connection_error
|
|
320
|
+
|
|
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
|
+
manager = MCPSessionManager.shared()
|
|
329
|
+
|
|
330
|
+
# All connection errors should use the simplified message format
|
|
331
|
+
with pytest.raises(RuntimeError, match="Unable to connect to MCP server"):
|
|
332
|
+
async with manager.mcp_client(tool_server):
|
|
333
|
+
pass
|
|
334
|
+
|
|
335
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
336
|
+
async def test_remote_mcp_http_error_in_nested_exceptions(self, mock_client):
|
|
337
|
+
"""Test remote MCP session extracts HTTP error from nested exceptions."""
|
|
338
|
+
# Create HTTP error nested in a mock exception group
|
|
339
|
+
response = MagicMock()
|
|
340
|
+
response.status_code = 401
|
|
341
|
+
response.reason_phrase = "Unauthorized"
|
|
342
|
+
http_error = httpx.HTTPStatusError(
|
|
343
|
+
"Unauthorized", request=MagicMock(), response=response
|
|
344
|
+
)
|
|
345
|
+
|
|
346
|
+
class MockExceptionGroup(Exception):
|
|
347
|
+
def __init__(self, exceptions):
|
|
348
|
+
super().__init__("Mock exception group")
|
|
349
|
+
self.exceptions = exceptions
|
|
350
|
+
|
|
351
|
+
group_error = MockExceptionGroup([ValueError("other error"), http_error])
|
|
352
|
+
|
|
353
|
+
# Mock client to raise the nested exception
|
|
354
|
+
mock_client.return_value.__aenter__.side_effect = group_error
|
|
355
|
+
|
|
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
|
+
manager = MCPSessionManager.shared()
|
|
364
|
+
|
|
365
|
+
# Should extract the HTTP error from the nested structure
|
|
366
|
+
with pytest.raises(
|
|
367
|
+
ValueError, match="The MCP server rejected the request. Status 401"
|
|
368
|
+
):
|
|
369
|
+
async with manager.mcp_client(tool_server):
|
|
370
|
+
pass
|
|
371
|
+
|
|
372
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
373
|
+
async def test_remote_mcp_connection_error_in_nested_exceptions(self, mock_client):
|
|
374
|
+
"""Test remote MCP session extracts connection error from nested exceptions."""
|
|
375
|
+
# Create connection error nested in mock exception group
|
|
376
|
+
connection_error = ConnectionError("Connection timeout")
|
|
377
|
+
|
|
378
|
+
class MockExceptionGroup(Exception):
|
|
379
|
+
def __init__(self, exceptions):
|
|
380
|
+
super().__init__("Mock exception group")
|
|
381
|
+
self.exceptions = exceptions
|
|
382
|
+
|
|
383
|
+
group_error = MockExceptionGroup([ValueError("other error"), connection_error])
|
|
384
|
+
|
|
385
|
+
# Mock client to raise the nested exception
|
|
386
|
+
mock_client.return_value.__aenter__.side_effect = group_error
|
|
387
|
+
|
|
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
|
+
manager = MCPSessionManager.shared()
|
|
396
|
+
|
|
397
|
+
# Should extract the connection error from the nested structure
|
|
398
|
+
with pytest.raises(RuntimeError, match="Unable to connect to MCP server"):
|
|
399
|
+
async with manager.mcp_client(tool_server):
|
|
400
|
+
pass
|
|
401
|
+
|
|
402
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
403
|
+
async def test_remote_mcp_unknown_error_fallback(self, mock_client):
|
|
404
|
+
"""Test remote MCP session handles unknown errors with fallback message."""
|
|
405
|
+
# Mock client to raise an unknown error type
|
|
406
|
+
unknown_error = RuntimeError("Unexpected error")
|
|
407
|
+
mock_client.return_value.__aenter__.side_effect = unknown_error
|
|
408
|
+
|
|
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
|
+
manager = MCPSessionManager.shared()
|
|
417
|
+
|
|
418
|
+
# Should use the fallback error message
|
|
419
|
+
with pytest.raises(RuntimeError, match="Failed to connect to the MCP Server"):
|
|
420
|
+
async with manager.mcp_client(tool_server):
|
|
421
|
+
pass
|
|
422
|
+
|
|
423
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
424
|
+
@patch("kiln_ai.utils.config.Config.shared")
|
|
425
|
+
async def test_session_with_secret_headers(self, mock_config, mock_client):
|
|
426
|
+
"""Test session creation with secret headers retrieved from config."""
|
|
427
|
+
# Mock the streams
|
|
428
|
+
mock_read_stream = MagicMock()
|
|
429
|
+
mock_write_stream = MagicMock()
|
|
430
|
+
|
|
431
|
+
# Configure the mock client context manager
|
|
432
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
433
|
+
mock_read_stream,
|
|
434
|
+
mock_write_stream,
|
|
435
|
+
None,
|
|
436
|
+
)
|
|
437
|
+
|
|
438
|
+
# Mock config with secret headers
|
|
439
|
+
mock_config_instance = MagicMock()
|
|
440
|
+
mock_config_instance.get_value.return_value = {
|
|
441
|
+
"test_server_id::Authorization": "Bearer secret-token-123",
|
|
442
|
+
"test_server_id::X-API-Key": "api-key-456",
|
|
443
|
+
"other_server::Token": "other-token", # Should be ignored
|
|
444
|
+
}
|
|
445
|
+
mock_config.return_value = mock_config_instance
|
|
446
|
+
|
|
447
|
+
# Create a tool server with secret header keys
|
|
448
|
+
tool_server = ExternalToolServer(
|
|
449
|
+
name="secret_headers_server",
|
|
450
|
+
type=ToolServerType.remote_mcp,
|
|
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"
|
|
460
|
+
|
|
461
|
+
manager = MCPSessionManager.shared()
|
|
462
|
+
|
|
463
|
+
with patch(
|
|
464
|
+
"kiln_ai.tools.mcp_session_manager.ClientSession"
|
|
465
|
+
) as mock_session_class:
|
|
466
|
+
mock_session_instance = AsyncMock()
|
|
467
|
+
mock_session_class.return_value.__aenter__.return_value = (
|
|
468
|
+
mock_session_instance
|
|
469
|
+
)
|
|
470
|
+
|
|
471
|
+
async with manager.mcp_client(tool_server) as session:
|
|
472
|
+
assert session is mock_session_instance
|
|
473
|
+
|
|
474
|
+
# Verify config was accessed for mcp_secrets
|
|
475
|
+
mock_config_instance.get_value.assert_called_once_with(MCP_SECRETS_KEY)
|
|
476
|
+
|
|
477
|
+
# Verify streamablehttp_client was called with merged headers
|
|
478
|
+
expected_headers = {
|
|
479
|
+
"Content-Type": "application/json",
|
|
480
|
+
"Authorization": "Bearer secret-token-123",
|
|
481
|
+
"X-API-Key": "api-key-456",
|
|
482
|
+
}
|
|
483
|
+
mock_client.assert_called_once_with(
|
|
484
|
+
"http://example.com/mcp", headers=expected_headers
|
|
485
|
+
)
|
|
486
|
+
|
|
487
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
488
|
+
@patch("kiln_ai.utils.config.Config.shared")
|
|
489
|
+
async def test_session_with_partial_secret_headers(self, mock_config, mock_client):
|
|
490
|
+
"""Test session creation when only some secret headers are found in config."""
|
|
491
|
+
# Mock the streams
|
|
492
|
+
mock_read_stream = MagicMock()
|
|
493
|
+
mock_write_stream = MagicMock()
|
|
494
|
+
|
|
495
|
+
# Configure the mock client context manager
|
|
496
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
497
|
+
mock_read_stream,
|
|
498
|
+
mock_write_stream,
|
|
499
|
+
None,
|
|
500
|
+
)
|
|
501
|
+
|
|
502
|
+
# Mock config with only one of the expected secret headers
|
|
503
|
+
mock_config_instance = MagicMock()
|
|
504
|
+
mock_config_instance.get_value.return_value = {
|
|
505
|
+
"test_server_id::Authorization": "Bearer found-token",
|
|
506
|
+
# Missing test_server_id::X-API-Key
|
|
507
|
+
}
|
|
508
|
+
mock_config.return_value = mock_config_instance
|
|
509
|
+
|
|
510
|
+
# Create a tool server expecting two secret headers
|
|
511
|
+
tool_server = ExternalToolServer(
|
|
512
|
+
name="partial_secret_server",
|
|
513
|
+
type=ToolServerType.remote_mcp,
|
|
514
|
+
description="Server with partial secret headers",
|
|
515
|
+
properties={
|
|
516
|
+
"server_url": "http://example.com/mcp",
|
|
517
|
+
"headers": {"Content-Type": "application/json"},
|
|
518
|
+
"secret_header_keys": ["Authorization", "X-API-Key"],
|
|
519
|
+
},
|
|
520
|
+
)
|
|
521
|
+
tool_server.id = "test_server_id"
|
|
522
|
+
|
|
523
|
+
manager = MCPSessionManager.shared()
|
|
524
|
+
|
|
525
|
+
with patch(
|
|
526
|
+
"kiln_ai.tools.mcp_session_manager.ClientSession"
|
|
527
|
+
) as mock_session_class:
|
|
528
|
+
mock_session_instance = AsyncMock()
|
|
529
|
+
mock_session_class.return_value.__aenter__.return_value = (
|
|
530
|
+
mock_session_instance
|
|
531
|
+
)
|
|
532
|
+
|
|
533
|
+
async with manager.mcp_client(tool_server) as session:
|
|
534
|
+
assert session is mock_session_instance
|
|
535
|
+
|
|
536
|
+
# Verify only the found secret header is merged
|
|
537
|
+
expected_headers = {
|
|
538
|
+
"Content-Type": "application/json",
|
|
539
|
+
"Authorization": "Bearer found-token",
|
|
540
|
+
# X-API-Key should not be present since it wasn't found in config
|
|
541
|
+
}
|
|
542
|
+
mock_client.assert_called_once_with(
|
|
543
|
+
"http://example.com/mcp", headers=expected_headers
|
|
544
|
+
)
|
|
545
|
+
|
|
546
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
547
|
+
@patch("kiln_ai.utils.config.Config.shared")
|
|
548
|
+
async def test_session_with_no_secret_headers_config(
|
|
549
|
+
self, mock_config, mock_client
|
|
550
|
+
):
|
|
551
|
+
"""Test session creation when config has no mcp_secrets."""
|
|
552
|
+
# Mock the streams
|
|
553
|
+
mock_read_stream = MagicMock()
|
|
554
|
+
mock_write_stream = MagicMock()
|
|
555
|
+
|
|
556
|
+
# Configure the mock client context manager
|
|
557
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
558
|
+
mock_read_stream,
|
|
559
|
+
mock_write_stream,
|
|
560
|
+
None,
|
|
561
|
+
)
|
|
562
|
+
|
|
563
|
+
# Mock config with no mcp_secrets (returns None)
|
|
564
|
+
mock_config_instance = MagicMock()
|
|
565
|
+
mock_config_instance.get_value.return_value = None
|
|
566
|
+
mock_config.return_value = mock_config_instance
|
|
567
|
+
|
|
568
|
+
# Create a tool server expecting secret headers
|
|
569
|
+
tool_server = ExternalToolServer(
|
|
570
|
+
name="no_secrets_config_server",
|
|
571
|
+
type=ToolServerType.remote_mcp,
|
|
572
|
+
description="Server with no secrets in config",
|
|
573
|
+
properties={
|
|
574
|
+
"server_url": "http://example.com/mcp",
|
|
575
|
+
"headers": {"Content-Type": "application/json"},
|
|
576
|
+
"secret_header_keys": ["Authorization"],
|
|
577
|
+
},
|
|
578
|
+
)
|
|
579
|
+
|
|
580
|
+
manager = MCPSessionManager.shared()
|
|
581
|
+
|
|
582
|
+
with patch(
|
|
583
|
+
"kiln_ai.tools.mcp_session_manager.ClientSession"
|
|
584
|
+
) as mock_session_class:
|
|
585
|
+
mock_session_instance = AsyncMock()
|
|
586
|
+
mock_session_class.return_value.__aenter__.return_value = (
|
|
587
|
+
mock_session_instance
|
|
588
|
+
)
|
|
589
|
+
|
|
590
|
+
async with manager.mcp_client(tool_server) as session:
|
|
591
|
+
assert session is mock_session_instance
|
|
592
|
+
|
|
593
|
+
# Verify only the original headers are used
|
|
594
|
+
expected_headers = {"Content-Type": "application/json"}
|
|
595
|
+
mock_client.assert_called_once_with(
|
|
596
|
+
"http://example.com/mcp", headers=expected_headers
|
|
597
|
+
)
|
|
598
|
+
|
|
599
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
600
|
+
async def test_session_with_empty_secret_header_keys(self, mock_client):
|
|
601
|
+
"""Test session creation with empty secret_header_keys list."""
|
|
602
|
+
# Mock the streams
|
|
603
|
+
mock_read_stream = MagicMock()
|
|
604
|
+
mock_write_stream = MagicMock()
|
|
605
|
+
|
|
606
|
+
# Configure the mock client context manager
|
|
607
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
608
|
+
mock_read_stream,
|
|
609
|
+
mock_write_stream,
|
|
610
|
+
None,
|
|
611
|
+
)
|
|
612
|
+
|
|
613
|
+
# Create a tool server with empty secret header keys
|
|
614
|
+
tool_server = ExternalToolServer(
|
|
615
|
+
name="empty_secret_keys_server",
|
|
616
|
+
type=ToolServerType.remote_mcp,
|
|
617
|
+
description="Server with empty secret header keys",
|
|
618
|
+
properties={
|
|
619
|
+
"server_url": "http://example.com/mcp",
|
|
620
|
+
"headers": {"Content-Type": "application/json"},
|
|
621
|
+
"secret_header_keys": [], # Empty list
|
|
622
|
+
},
|
|
623
|
+
)
|
|
624
|
+
|
|
625
|
+
manager = MCPSessionManager.shared()
|
|
626
|
+
|
|
627
|
+
with patch(
|
|
628
|
+
"kiln_ai.tools.mcp_session_manager.ClientSession"
|
|
629
|
+
) as mock_session_class:
|
|
630
|
+
mock_session_instance = AsyncMock()
|
|
631
|
+
mock_session_class.return_value.__aenter__.return_value = (
|
|
632
|
+
mock_session_instance
|
|
633
|
+
)
|
|
634
|
+
|
|
635
|
+
async with manager.mcp_client(tool_server) as session:
|
|
636
|
+
assert session is mock_session_instance
|
|
637
|
+
|
|
638
|
+
# Verify only the original headers are used (no config access needed for empty list)
|
|
639
|
+
expected_headers = {"Content-Type": "application/json"}
|
|
640
|
+
mock_client.assert_called_once_with(
|
|
641
|
+
"http://example.com/mcp", headers=expected_headers
|
|
642
|
+
)
|
|
643
|
+
|
|
644
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
645
|
+
async def test_session_with_missing_secret_header_keys_property(self, mock_client):
|
|
646
|
+
"""Test session creation when secret_header_keys property is missing."""
|
|
647
|
+
# Mock the streams
|
|
648
|
+
mock_read_stream = MagicMock()
|
|
649
|
+
mock_write_stream = MagicMock()
|
|
650
|
+
|
|
651
|
+
# Configure the mock client context manager
|
|
652
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
653
|
+
mock_read_stream,
|
|
654
|
+
mock_write_stream,
|
|
655
|
+
None,
|
|
656
|
+
)
|
|
657
|
+
|
|
658
|
+
# Create a tool server without secret_header_keys property
|
|
659
|
+
tool_server = ExternalToolServer(
|
|
660
|
+
name="missing_secret_keys_server",
|
|
661
|
+
type=ToolServerType.remote_mcp,
|
|
662
|
+
description="Server without secret header keys property",
|
|
663
|
+
properties={
|
|
664
|
+
"server_url": "http://example.com/mcp",
|
|
665
|
+
"headers": {"Content-Type": "application/json"},
|
|
666
|
+
# No secret_header_keys property
|
|
667
|
+
},
|
|
668
|
+
)
|
|
669
|
+
|
|
670
|
+
manager = MCPSessionManager.shared()
|
|
671
|
+
|
|
672
|
+
with patch(
|
|
673
|
+
"kiln_ai.tools.mcp_session_manager.ClientSession"
|
|
674
|
+
) as mock_session_class:
|
|
675
|
+
mock_session_instance = AsyncMock()
|
|
676
|
+
mock_session_class.return_value.__aenter__.return_value = (
|
|
677
|
+
mock_session_instance
|
|
678
|
+
)
|
|
679
|
+
|
|
680
|
+
async with manager.mcp_client(tool_server) as session:
|
|
681
|
+
assert session is mock_session_instance
|
|
682
|
+
|
|
683
|
+
# Verify only the original headers are used (no config access needed when property missing)
|
|
684
|
+
expected_headers = {"Content-Type": "application/json"}
|
|
685
|
+
mock_client.assert_called_once_with(
|
|
686
|
+
"http://example.com/mcp", headers=expected_headers
|
|
687
|
+
)
|
|
688
|
+
|
|
689
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
690
|
+
@patch("kiln_ai.utils.config.Config.shared")
|
|
691
|
+
async def test_secret_headers_do_not_modify_original_properties(
|
|
692
|
+
self, mock_config, mock_client
|
|
693
|
+
):
|
|
694
|
+
"""Test that secret headers are not saved back to the original tool server properties."""
|
|
695
|
+
# Mock the streams
|
|
696
|
+
mock_read_stream = MagicMock()
|
|
697
|
+
mock_write_stream = MagicMock()
|
|
698
|
+
|
|
699
|
+
# Configure the mock client context manager
|
|
700
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
701
|
+
mock_read_stream,
|
|
702
|
+
mock_write_stream,
|
|
703
|
+
None,
|
|
704
|
+
)
|
|
705
|
+
|
|
706
|
+
# Mock config with secret headers
|
|
707
|
+
mock_config_instance = MagicMock()
|
|
708
|
+
mock_config_instance.get_value.return_value = {
|
|
709
|
+
"test_server_id::Authorization": "Bearer secret-token-123",
|
|
710
|
+
"test_server_id::X-API-Key": "api-key-456",
|
|
711
|
+
}
|
|
712
|
+
mock_config.return_value = mock_config_instance
|
|
713
|
+
|
|
714
|
+
# Create a tool server with secret header keys
|
|
715
|
+
tool_server = ExternalToolServer(
|
|
716
|
+
name="bug_test_server",
|
|
717
|
+
type=ToolServerType.remote_mcp,
|
|
718
|
+
description="Server to test the secret headers bug",
|
|
719
|
+
properties={
|
|
720
|
+
"server_url": "http://example.com/mcp",
|
|
721
|
+
"headers": {"Content-Type": "application/json"},
|
|
722
|
+
"secret_header_keys": ["Authorization", "X-API-Key"],
|
|
723
|
+
},
|
|
724
|
+
)
|
|
725
|
+
# Set the server ID to match our mock secrets
|
|
726
|
+
tool_server.id = "test_server_id"
|
|
727
|
+
|
|
728
|
+
# Store original headers for comparison
|
|
729
|
+
original_headers = tool_server.properties["headers"].copy()
|
|
730
|
+
|
|
731
|
+
manager = MCPSessionManager.shared()
|
|
732
|
+
|
|
733
|
+
with patch(
|
|
734
|
+
"kiln_ai.tools.mcp_session_manager.ClientSession"
|
|
735
|
+
) as mock_session_class:
|
|
736
|
+
mock_session_instance = AsyncMock()
|
|
737
|
+
mock_session_class.return_value.__aenter__.return_value = (
|
|
738
|
+
mock_session_instance
|
|
739
|
+
)
|
|
740
|
+
|
|
741
|
+
# Use the session multiple times to ensure the bug doesn't occur
|
|
742
|
+
async with manager.mcp_client(tool_server) as session:
|
|
743
|
+
assert session is mock_session_instance
|
|
744
|
+
|
|
745
|
+
# Check that original headers are unchanged after first use
|
|
746
|
+
assert tool_server.properties["headers"] == original_headers
|
|
747
|
+
assert "Authorization" not in tool_server.properties["headers"]
|
|
748
|
+
assert "X-API-Key" not in tool_server.properties["headers"]
|
|
749
|
+
|
|
750
|
+
# Use the session a second time to ensure the bug doesn't occur on subsequent uses
|
|
751
|
+
async with manager.mcp_client(tool_server) as session:
|
|
752
|
+
assert session is mock_session_instance
|
|
753
|
+
|
|
754
|
+
# Check that original headers are still unchanged after second use
|
|
755
|
+
assert tool_server.properties["headers"] == original_headers
|
|
756
|
+
assert "Authorization" not in tool_server.properties["headers"]
|
|
757
|
+
assert "X-API-Key" not in tool_server.properties["headers"]
|
|
758
|
+
|
|
759
|
+
# Verify streamablehttp_client was called with merged headers both times
|
|
760
|
+
expected_headers = {
|
|
761
|
+
"Content-Type": "application/json",
|
|
762
|
+
"Authorization": "Bearer secret-token-123",
|
|
763
|
+
"X-API-Key": "api-key-456",
|
|
764
|
+
}
|
|
765
|
+
# Should have been called twice (once for each session)
|
|
766
|
+
assert mock_client.call_count == 2
|
|
767
|
+
for call in mock_client.call_args_list:
|
|
768
|
+
assert call[0][0] == "http://example.com/mcp"
|
|
769
|
+
assert call[1]["headers"] == expected_headers
|
|
770
|
+
|
|
771
|
+
@patch("kiln_ai.tools.mcp_session_manager.streamablehttp_client")
|
|
772
|
+
@patch("kiln_ai.utils.config.Config.shared")
|
|
773
|
+
async def test_demonstrates_bug_without_copy_fix(self, mock_config, mock_client):
|
|
774
|
+
"""
|
|
775
|
+
Test that demonstrates the bug that would occur without the .copy() fix.
|
|
776
|
+
This test simulates what would happen if we modified headers directly.
|
|
777
|
+
"""
|
|
778
|
+
# Mock the streams
|
|
779
|
+
mock_read_stream = MagicMock()
|
|
780
|
+
mock_write_stream = MagicMock()
|
|
781
|
+
|
|
782
|
+
# Configure the mock client context manager
|
|
783
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
784
|
+
mock_read_stream,
|
|
785
|
+
mock_write_stream,
|
|
786
|
+
None,
|
|
787
|
+
)
|
|
788
|
+
|
|
789
|
+
# Mock config with secret headers
|
|
790
|
+
mock_config_instance = MagicMock()
|
|
791
|
+
mock_config_instance.get_value.return_value = {
|
|
792
|
+
"test_server_id::Authorization": "Bearer secret-token-123",
|
|
793
|
+
}
|
|
794
|
+
mock_config.return_value = mock_config_instance
|
|
795
|
+
|
|
796
|
+
# Create a tool server with secret header keys
|
|
797
|
+
tool_server = ExternalToolServer(
|
|
798
|
+
name="bug_demo_server",
|
|
799
|
+
type=ToolServerType.remote_mcp,
|
|
800
|
+
description="Server to demonstrate the bug",
|
|
801
|
+
properties={
|
|
802
|
+
"server_url": "http://example.com/mcp",
|
|
803
|
+
"headers": {"Content-Type": "application/json"},
|
|
804
|
+
"secret_header_keys": ["Authorization"],
|
|
805
|
+
},
|
|
806
|
+
)
|
|
807
|
+
tool_server.id = "test_server_id"
|
|
808
|
+
|
|
809
|
+
# Store original headers for comparison
|
|
810
|
+
original_headers = tool_server.properties["headers"].copy()
|
|
811
|
+
|
|
812
|
+
# Simulate the buggy behavior by directly modifying the headers
|
|
813
|
+
# (This is what would happen without the .copy() fix)
|
|
814
|
+
buggy_headers = tool_server.properties.get("headers", {}) # No .copy()!
|
|
815
|
+
|
|
816
|
+
# Simulate what the buggy code would do - directly modify the original headers
|
|
817
|
+
secret_headers_keys = tool_server.properties.get("secret_header_keys", [])
|
|
818
|
+
if secret_headers_keys:
|
|
819
|
+
config = mock_config_instance
|
|
820
|
+
mcp_secrets = config.get_value(MCP_SECRETS_KEY)
|
|
821
|
+
if mcp_secrets:
|
|
822
|
+
for header_name in secret_headers_keys:
|
|
823
|
+
header_value = mcp_secrets.get(f"{tool_server.id}::{header_name}")
|
|
824
|
+
if header_value:
|
|
825
|
+
buggy_headers[header_name] = header_value
|
|
826
|
+
|
|
827
|
+
# Now the original properties would be contaminated with secrets!
|
|
828
|
+
assert "Authorization" in tool_server.properties["headers"]
|
|
829
|
+
assert (
|
|
830
|
+
tool_server.properties["headers"]["Authorization"]
|
|
831
|
+
== "Bearer secret-token-123"
|
|
832
|
+
)
|
|
833
|
+
|
|
834
|
+
# This demonstrates the security bug - secrets are now permanently stored
|
|
835
|
+
# in the tool server properties and would be serialized/saved
|
|
836
|
+
assert tool_server.properties["headers"] != original_headers
|
|
837
|
+
|
|
838
|
+
@patch("kiln_ai.tools.mcp_session_manager.stdio_client")
|
|
839
|
+
async def test_local_mcp_session_creation(self, mock_client):
|
|
840
|
+
"""Test successful local MCP session creation with mocked client."""
|
|
841
|
+
# Mock the streams
|
|
842
|
+
mock_read_stream = MagicMock()
|
|
843
|
+
mock_write_stream = MagicMock()
|
|
844
|
+
|
|
845
|
+
# Configure the mock client context manager
|
|
846
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
847
|
+
mock_read_stream,
|
|
848
|
+
mock_write_stream,
|
|
849
|
+
)
|
|
850
|
+
|
|
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
|
+
manager = MCPSessionManager.shared()
|
|
864
|
+
|
|
865
|
+
with patch(
|
|
866
|
+
"kiln_ai.tools.mcp_session_manager.ClientSession"
|
|
867
|
+
) as mock_session_class:
|
|
868
|
+
mock_session_instance = AsyncMock()
|
|
869
|
+
mock_session_class.return_value.__aenter__.return_value = (
|
|
870
|
+
mock_session_instance
|
|
871
|
+
)
|
|
872
|
+
|
|
873
|
+
async with manager.mcp_client(tool_server) as session:
|
|
874
|
+
# Verify session is returned
|
|
875
|
+
assert session is mock_session_instance
|
|
876
|
+
|
|
877
|
+
# Verify initialize was called
|
|
878
|
+
mock_session_instance.initialize.assert_called_once()
|
|
879
|
+
|
|
880
|
+
# Verify stdio_client was called with correct parameters
|
|
881
|
+
call_args = mock_client.call_args[0][0] # Get the StdioServerParameters
|
|
882
|
+
assert call_args.command == "python"
|
|
883
|
+
assert call_args.args == ["-m", "my_mcp_server"]
|
|
884
|
+
# Verify that the original env vars are included plus PATH
|
|
885
|
+
assert "API_KEY" in call_args.env
|
|
886
|
+
assert call_args.env["API_KEY"] == "test123"
|
|
887
|
+
assert "PATH" in call_args.env
|
|
888
|
+
assert len(call_args.env["PATH"]) > 0
|
|
889
|
+
|
|
890
|
+
@patch("kiln_ai.tools.mcp_session_manager.stdio_client")
|
|
891
|
+
async def test_local_mcp_session_with_defaults(self, mock_client):
|
|
892
|
+
"""Test local MCP session creation with default env_vars."""
|
|
893
|
+
# Mock the streams
|
|
894
|
+
mock_read_stream = MagicMock()
|
|
895
|
+
mock_write_stream = MagicMock()
|
|
896
|
+
|
|
897
|
+
# Configure the mock client context manager
|
|
898
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
899
|
+
mock_read_stream,
|
|
900
|
+
mock_write_stream,
|
|
901
|
+
)
|
|
902
|
+
|
|
903
|
+
# Create a tool server without env_vars (should default to {})
|
|
904
|
+
tool_server = ExternalToolServer(
|
|
905
|
+
name="test_local_server_defaults",
|
|
906
|
+
type=ToolServerType.local_mcp,
|
|
907
|
+
description="Test local server with defaults",
|
|
908
|
+
properties={
|
|
909
|
+
"command": "node",
|
|
910
|
+
"args": ["server.js"],
|
|
911
|
+
# No env_vars provided
|
|
912
|
+
},
|
|
913
|
+
)
|
|
914
|
+
|
|
915
|
+
manager = MCPSessionManager.shared()
|
|
916
|
+
|
|
917
|
+
with patch(
|
|
918
|
+
"kiln_ai.tools.mcp_session_manager.ClientSession"
|
|
919
|
+
) as mock_session_class:
|
|
920
|
+
mock_session_instance = AsyncMock()
|
|
921
|
+
mock_session_class.return_value.__aenter__.return_value = (
|
|
922
|
+
mock_session_instance
|
|
923
|
+
)
|
|
924
|
+
|
|
925
|
+
async with manager.mcp_client(tool_server) as session:
|
|
926
|
+
assert session is mock_session_instance
|
|
927
|
+
|
|
928
|
+
# Verify stdio_client was called with PATH automatically added
|
|
929
|
+
call_args = mock_client.call_args[0][0]
|
|
930
|
+
# Should only contain PATH (no other env vars were provided)
|
|
931
|
+
assert "PATH" in call_args.env
|
|
932
|
+
assert len(call_args.env["PATH"]) > 0
|
|
933
|
+
# Should not have any other env vars besides PATH
|
|
934
|
+
assert len(call_args.env) == 1
|
|
935
|
+
|
|
936
|
+
async def test_local_mcp_missing_command_error(self):
|
|
937
|
+
"""Test that missing command raises ValueError for local MCP."""
|
|
938
|
+
with pytest.raises(
|
|
939
|
+
ValidationError,
|
|
940
|
+
match="command must be a string to start a local MCP server",
|
|
941
|
+
):
|
|
942
|
+
ExternalToolServer(
|
|
943
|
+
name="missing_command_server",
|
|
944
|
+
type=ToolServerType.local_mcp,
|
|
945
|
+
description="Server missing command",
|
|
946
|
+
properties={
|
|
947
|
+
# No command provided
|
|
948
|
+
"args": ["arg1"],
|
|
949
|
+
"env_vars": {},
|
|
950
|
+
},
|
|
951
|
+
)
|
|
952
|
+
|
|
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
|
+
async def test_local_mcp_empty_args_allowed(self):
|
|
971
|
+
"""Test that empty args list is now allowed for local MCP."""
|
|
972
|
+
# Should not raise any exception - empty args are now allowed
|
|
973
|
+
tool_server = ExternalToolServer(
|
|
974
|
+
name="empty_args_server",
|
|
975
|
+
type=ToolServerType.local_mcp,
|
|
976
|
+
description="Server with empty args",
|
|
977
|
+
properties={
|
|
978
|
+
"command": "python",
|
|
979
|
+
"args": [], # Empty args list should now be allowed
|
|
980
|
+
"env_vars": {},
|
|
981
|
+
},
|
|
982
|
+
)
|
|
983
|
+
|
|
984
|
+
assert tool_server.name == "empty_args_server"
|
|
985
|
+
assert tool_server.type == ToolServerType.local_mcp
|
|
986
|
+
assert tool_server.properties["args"] == []
|
|
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
|
|
1039
|
+
|
|
1040
|
+
@patch("kiln_ai.tools.mcp_session_manager.stdio_client")
|
|
1041
|
+
@patch("kiln_ai.utils.config.Config.shared")
|
|
1042
|
+
async def test_local_mcp_session_with_secrets(self, mock_config, mock_client):
|
|
1043
|
+
"""Test local MCP session creation with secret environment variables."""
|
|
1044
|
+
# Mock config to return different values based on the key
|
|
1045
|
+
mock_config_instance = MagicMock()
|
|
1046
|
+
|
|
1047
|
+
def mock_get_value(key):
|
|
1048
|
+
if key == MCP_SECRETS_KEY:
|
|
1049
|
+
return {
|
|
1050
|
+
"test_server_id::SECRET_API_KEY": "secret_value_123",
|
|
1051
|
+
"test_server_id::ANOTHER_SECRET": "another_secret_value",
|
|
1052
|
+
}
|
|
1053
|
+
elif key == "custom_mcp_path":
|
|
1054
|
+
return None # No custom path, will use shell path
|
|
1055
|
+
return None
|
|
1056
|
+
|
|
1057
|
+
mock_config_instance.get_value.side_effect = mock_get_value
|
|
1058
|
+
mock_config.return_value = mock_config_instance
|
|
1059
|
+
|
|
1060
|
+
# Mock the streams
|
|
1061
|
+
mock_read_stream = MagicMock()
|
|
1062
|
+
mock_write_stream = MagicMock()
|
|
1063
|
+
|
|
1064
|
+
# Configure the mock client context manager
|
|
1065
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
1066
|
+
mock_read_stream,
|
|
1067
|
+
mock_write_stream,
|
|
1068
|
+
)
|
|
1069
|
+
|
|
1070
|
+
# Create a tool server with secret env var keys
|
|
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
|
|
1083
|
+
tool_server.id = "test_server_id"
|
|
1084
|
+
|
|
1085
|
+
manager = MCPSessionManager.shared()
|
|
1086
|
+
|
|
1087
|
+
# Mock get_shell_path to return a simple PATH
|
|
1088
|
+
with (
|
|
1089
|
+
patch.object(manager, "get_shell_path", return_value="/usr/bin:/bin"),
|
|
1090
|
+
patch(
|
|
1091
|
+
"kiln_ai.tools.mcp_session_manager.ClientSession"
|
|
1092
|
+
) as mock_session_class,
|
|
1093
|
+
):
|
|
1094
|
+
mock_session_instance = AsyncMock()
|
|
1095
|
+
mock_session_class.return_value.__aenter__.return_value = (
|
|
1096
|
+
mock_session_instance
|
|
1097
|
+
)
|
|
1098
|
+
|
|
1099
|
+
async with manager.mcp_client(tool_server) as session:
|
|
1100
|
+
# Verify session is returned
|
|
1101
|
+
assert session is mock_session_instance
|
|
1102
|
+
|
|
1103
|
+
# Verify initialize was called
|
|
1104
|
+
mock_session_instance.initialize.assert_called_once()
|
|
1105
|
+
|
|
1106
|
+
# Verify config was accessed for mcp_secrets
|
|
1107
|
+
assert mock_config_instance.get_value.call_count == 2
|
|
1108
|
+
mock_config_instance.get_value.assert_any_call(MCP_SECRETS_KEY)
|
|
1109
|
+
mock_config_instance.get_value.assert_any_call("custom_mcp_path")
|
|
1110
|
+
|
|
1111
|
+
# Verify stdio_client was called with correct parameters including secrets
|
|
1112
|
+
call_args = mock_client.call_args[0][0] # Get the StdioServerParameters
|
|
1113
|
+
assert call_args.command == "python"
|
|
1114
|
+
assert call_args.args == ["-m", "my_server"]
|
|
1115
|
+
|
|
1116
|
+
# Verify that both public and secret env vars are included
|
|
1117
|
+
assert "PUBLIC_VAR" in call_args.env
|
|
1118
|
+
assert call_args.env["PUBLIC_VAR"] == "public_value"
|
|
1119
|
+
assert "SECRET_API_KEY" in call_args.env
|
|
1120
|
+
assert call_args.env["SECRET_API_KEY"] == "secret_value_123"
|
|
1121
|
+
assert "ANOTHER_SECRET" in call_args.env
|
|
1122
|
+
assert call_args.env["ANOTHER_SECRET"] == "another_secret_value"
|
|
1123
|
+
assert "PATH" in call_args.env
|
|
1124
|
+
|
|
1125
|
+
# Verify original properties were not modified (security check)
|
|
1126
|
+
original_env_vars = tool_server.properties.get("env_vars", {})
|
|
1127
|
+
assert "SECRET_API_KEY" not in original_env_vars
|
|
1128
|
+
assert "ANOTHER_SECRET" not in original_env_vars
|
|
1129
|
+
assert original_env_vars.get("PUBLIC_VAR") == "public_value"
|
|
1130
|
+
|
|
1131
|
+
@pytest.mark.parametrize(
|
|
1132
|
+
"error_type,error_message",
|
|
1133
|
+
[
|
|
1134
|
+
(McpError, "MCP initialization failed"),
|
|
1135
|
+
(FileNotFoundError, "Command 'nonexistent' not found"),
|
|
1136
|
+
(RuntimeError, "Unknown server error"),
|
|
1137
|
+
(ValueError, "Invalid arguments provided"),
|
|
1138
|
+
],
|
|
1139
|
+
)
|
|
1140
|
+
@patch("kiln_ai.tools.mcp_session_manager.stdio_client")
|
|
1141
|
+
async def test_local_mcp_various_errors_use_simplified_message(
|
|
1142
|
+
self, mock_client, error_type, error_message
|
|
1143
|
+
):
|
|
1144
|
+
"""Test local MCP session handles various errors with simplified message."""
|
|
1145
|
+
# Create the appropriate error
|
|
1146
|
+
if error_type == McpError:
|
|
1147
|
+
error_data = ErrorData(code=-1, message=error_message)
|
|
1148
|
+
test_error = McpError(error_data)
|
|
1149
|
+
else:
|
|
1150
|
+
test_error = error_type(error_message)
|
|
1151
|
+
|
|
1152
|
+
# Mock client to raise the error
|
|
1153
|
+
mock_client.return_value.__aenter__.side_effect = test_error
|
|
1154
|
+
|
|
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
|
+
manager = MCPSessionManager.shared()
|
|
1167
|
+
|
|
1168
|
+
# All local errors should now use the simplified message format
|
|
1169
|
+
with pytest.raises(RuntimeError, match="MCP server failed to start"):
|
|
1170
|
+
async with manager.mcp_client(tool_server):
|
|
1171
|
+
pass
|
|
1172
|
+
|
|
1173
|
+
@patch("kiln_ai.tools.mcp_session_manager.stdio_client")
|
|
1174
|
+
async def test_local_mcp_mcp_error_in_nested_exceptions(self, mock_client):
|
|
1175
|
+
"""Test local MCP session extracts McpError from nested exceptions."""
|
|
1176
|
+
# Create McpError nested in mock exception group
|
|
1177
|
+
error_data = ErrorData(code=-1, message="Server startup failed")
|
|
1178
|
+
mcp_error = McpError(error_data)
|
|
1179
|
+
|
|
1180
|
+
class MockExceptionGroup(Exception):
|
|
1181
|
+
def __init__(self, exceptions):
|
|
1182
|
+
super().__init__("Mock exception group")
|
|
1183
|
+
self.exceptions = exceptions
|
|
1184
|
+
|
|
1185
|
+
group_error = MockExceptionGroup([ValueError("other error"), mcp_error])
|
|
1186
|
+
|
|
1187
|
+
# Mock client to raise the nested exception
|
|
1188
|
+
mock_client.return_value.__aenter__.side_effect = group_error
|
|
1189
|
+
|
|
1190
|
+
tool_server = ExternalToolServer(
|
|
1191
|
+
name="test_server",
|
|
1192
|
+
type=ToolServerType.local_mcp,
|
|
1193
|
+
description="Test server",
|
|
1194
|
+
properties={
|
|
1195
|
+
"command": "python",
|
|
1196
|
+
"args": ["-m", "broken_server"],
|
|
1197
|
+
"env_vars": {},
|
|
1198
|
+
},
|
|
1199
|
+
)
|
|
1200
|
+
|
|
1201
|
+
manager = MCPSessionManager.shared()
|
|
1202
|
+
|
|
1203
|
+
# Should extract the McpError from the nested structure and use simplified message
|
|
1204
|
+
with pytest.raises(RuntimeError, match="MCP server failed to start"):
|
|
1205
|
+
async with manager.mcp_client(tool_server):
|
|
1206
|
+
pass
|
|
1207
|
+
|
|
1208
|
+
def test_raise_local_mcp_error_method(self):
|
|
1209
|
+
"""Test the _raise_local_mcp_error helper method."""
|
|
1210
|
+
manager = MCPSessionManager()
|
|
1211
|
+
|
|
1212
|
+
# Test with different exception types
|
|
1213
|
+
test_exceptions = [
|
|
1214
|
+
ValueError("test value error"),
|
|
1215
|
+
FileNotFoundError("file not found"),
|
|
1216
|
+
RuntimeError("runtime error"),
|
|
1217
|
+
Exception("generic exception"),
|
|
1218
|
+
]
|
|
1219
|
+
|
|
1220
|
+
for original_error in test_exceptions:
|
|
1221
|
+
with pytest.raises(RuntimeError) as exc_info:
|
|
1222
|
+
manager._raise_local_mcp_error(original_error)
|
|
1223
|
+
|
|
1224
|
+
# Check that the error message contains expected text
|
|
1225
|
+
assert "MCP server failed to start" in str(exc_info.value)
|
|
1226
|
+
assert (
|
|
1227
|
+
"Please verify your command, arguments, and environment variables"
|
|
1228
|
+
in str(exc_info.value)
|
|
1229
|
+
)
|
|
1230
|
+
assert str(original_error) in str(exc_info.value)
|
|
1231
|
+
|
|
1232
|
+
# Check that the original exception is chained
|
|
1233
|
+
assert exc_info.value.__cause__ is original_error
|
|
1234
|
+
|
|
1235
|
+
@patch("kiln_ai.tools.mcp_session_manager.stdio_client")
|
|
1236
|
+
@patch("kiln_ai.utils.config.Config.shared")
|
|
1237
|
+
async def test_local_mcp_session_with_no_secrets_config(
|
|
1238
|
+
self, mock_config, mock_client
|
|
1239
|
+
):
|
|
1240
|
+
"""Test local MCP session creation when config has no mcp_secrets."""
|
|
1241
|
+
# Mock config to return None for mcp_secrets
|
|
1242
|
+
mock_config_instance = MagicMock()
|
|
1243
|
+
|
|
1244
|
+
def mock_get_value(key):
|
|
1245
|
+
if key == MCP_SECRETS_KEY:
|
|
1246
|
+
return None
|
|
1247
|
+
elif key == "custom_mcp_path":
|
|
1248
|
+
return None # No custom path, will use shell path
|
|
1249
|
+
return None
|
|
1250
|
+
|
|
1251
|
+
mock_config_instance.get_value.side_effect = mock_get_value
|
|
1252
|
+
mock_config.return_value = mock_config_instance
|
|
1253
|
+
|
|
1254
|
+
# Mock the streams
|
|
1255
|
+
mock_read_stream = MagicMock()
|
|
1256
|
+
mock_write_stream = MagicMock()
|
|
1257
|
+
mock_client.return_value.__aenter__.return_value = (
|
|
1258
|
+
mock_read_stream,
|
|
1259
|
+
mock_write_stream,
|
|
1260
|
+
)
|
|
1261
|
+
|
|
1262
|
+
# Create a tool server with secret env var keys but no secrets in config
|
|
1263
|
+
tool_server = ExternalToolServer(
|
|
1264
|
+
name="no_secrets_config_server",
|
|
1265
|
+
type=ToolServerType.local_mcp,
|
|
1266
|
+
description="Server with no secrets in config",
|
|
1267
|
+
properties={
|
|
1268
|
+
"command": "python",
|
|
1269
|
+
"args": ["-m", "my_server"],
|
|
1270
|
+
"env_vars": {"PUBLIC_VAR": "public_value"},
|
|
1271
|
+
"secret_env_var_keys": ["SECRET_API_KEY"],
|
|
1272
|
+
},
|
|
1273
|
+
)
|
|
1274
|
+
tool_server.id = "test_server_id"
|
|
1275
|
+
|
|
1276
|
+
manager = MCPSessionManager.shared()
|
|
1277
|
+
|
|
1278
|
+
# Mock get_shell_path to return a simple PATH
|
|
1279
|
+
with (
|
|
1280
|
+
patch.object(manager, "get_shell_path", return_value="/usr/bin:/bin"),
|
|
1281
|
+
patch(
|
|
1282
|
+
"kiln_ai.tools.mcp_session_manager.ClientSession"
|
|
1283
|
+
) as mock_session_class,
|
|
1284
|
+
):
|
|
1285
|
+
mock_session_instance = AsyncMock()
|
|
1286
|
+
mock_session_class.return_value.__aenter__.return_value = (
|
|
1287
|
+
mock_session_instance
|
|
1288
|
+
)
|
|
1289
|
+
|
|
1290
|
+
async with manager.mcp_client(tool_server):
|
|
1291
|
+
pass # Should not raise any errors
|
|
1292
|
+
|
|
1293
|
+
# Verify stdio_client was called and only public vars are included
|
|
1294
|
+
call_args = mock_client.call_args[0][0]
|
|
1295
|
+
assert "PUBLIC_VAR" in call_args.env
|
|
1296
|
+
assert call_args.env["PUBLIC_VAR"] == "public_value"
|
|
1297
|
+
assert "SECRET_API_KEY" not in call_args.env # Secret not found in config
|
|
1298
|
+
assert "PATH" in call_args.env
|
|
1299
|
+
|
|
1300
|
+
@patch("kiln_ai.utils.config.Config.shared")
|
|
1301
|
+
def test_get_path_with_custom_mcp_path(self, mock_config):
|
|
1302
|
+
"""Test _get_path() returns custom MCP path when configured."""
|
|
1303
|
+
# Setup mock config to return a custom path
|
|
1304
|
+
mock_config_instance = MagicMock()
|
|
1305
|
+
mock_config_instance.get_value.return_value = "/custom/mcp/path"
|
|
1306
|
+
mock_config.return_value = mock_config_instance
|
|
1307
|
+
|
|
1308
|
+
manager = MCPSessionManager()
|
|
1309
|
+
|
|
1310
|
+
# Mock get_shell_path to ensure it's not called
|
|
1311
|
+
with patch.object(manager, "get_shell_path") as mock_get_shell_path:
|
|
1312
|
+
result = manager._get_path()
|
|
1313
|
+
|
|
1314
|
+
assert result == "/custom/mcp/path"
|
|
1315
|
+
mock_config_instance.get_value.assert_called_once_with("custom_mcp_path")
|
|
1316
|
+
mock_get_shell_path.assert_not_called()
|
|
1317
|
+
|
|
1318
|
+
@patch("kiln_ai.utils.config.Config.shared")
|
|
1319
|
+
def test_get_path_fallback_to_shell_path(self, mock_config):
|
|
1320
|
+
"""Test _get_path() falls back to get_shell_path() when no custom path."""
|
|
1321
|
+
# Setup mock config to return None (no custom path)
|
|
1322
|
+
mock_config_instance = MagicMock()
|
|
1323
|
+
mock_config_instance.get_value.return_value = None
|
|
1324
|
+
mock_config.return_value = mock_config_instance
|
|
1325
|
+
|
|
1326
|
+
manager = MCPSessionManager()
|
|
1327
|
+
|
|
1328
|
+
with patch.object(
|
|
1329
|
+
manager, "get_shell_path", return_value="/shell/path"
|
|
1330
|
+
) as mock_shell:
|
|
1331
|
+
result = manager._get_path()
|
|
1332
|
+
|
|
1333
|
+
assert result == "/shell/path"
|
|
1334
|
+
mock_shell.assert_called_once()
|
|
1335
|
+
|
|
1336
|
+
@patch("sys.platform", "win32")
|
|
1337
|
+
@patch.dict(os.environ, {"PATH": "/windows/path"})
|
|
1338
|
+
def test_get_shell_path_windows(self):
|
|
1339
|
+
"""Test get_shell_path() on Windows platform."""
|
|
1340
|
+
manager = MCPSessionManager()
|
|
1341
|
+
|
|
1342
|
+
result = manager.get_shell_path()
|
|
1343
|
+
|
|
1344
|
+
assert result == "/windows/path"
|
|
1345
|
+
|
|
1346
|
+
@patch("sys.platform", "Windows")
|
|
1347
|
+
@patch.dict(os.environ, {"PATH": "/windows/path2"})
|
|
1348
|
+
def test_get_shell_path_windows_alt_platform_name(self):
|
|
1349
|
+
"""Test get_shell_path() on Windows with 'Windows' platform name."""
|
|
1350
|
+
manager = MCPSessionManager()
|
|
1351
|
+
|
|
1352
|
+
result = manager.get_shell_path()
|
|
1353
|
+
|
|
1354
|
+
assert result == "/windows/path2"
|
|
1355
|
+
|
|
1356
|
+
@patch("sys.platform", "linux")
|
|
1357
|
+
@patch.dict(os.environ, {"SHELL": "/bin/bash", "PATH": "/fallback/path"})
|
|
1358
|
+
@patch("subprocess.run")
|
|
1359
|
+
def test_get_shell_path_unix_success(self, mock_run):
|
|
1360
|
+
"""Test get_shell_path() successful shell execution on Unix."""
|
|
1361
|
+
# Mock successful subprocess execution
|
|
1362
|
+
mock_result = MagicMock()
|
|
1363
|
+
mock_result.returncode = 0
|
|
1364
|
+
mock_result.stdout = "/usr/local/bin:/usr/bin:/bin\n"
|
|
1365
|
+
mock_run.return_value = mock_result
|
|
1366
|
+
|
|
1367
|
+
manager = MCPSessionManager()
|
|
1368
|
+
|
|
1369
|
+
result = manager.get_shell_path()
|
|
1370
|
+
|
|
1371
|
+
assert result == "/usr/local/bin:/usr/bin:/bin"
|
|
1372
|
+
mock_run.assert_called_once_with(
|
|
1373
|
+
["/bin/bash", "-l", "-c", "echo $PATH"],
|
|
1374
|
+
capture_output=True,
|
|
1375
|
+
text=True,
|
|
1376
|
+
timeout=3,
|
|
1377
|
+
)
|
|
1378
|
+
|
|
1379
|
+
@patch("sys.platform", "linux")
|
|
1380
|
+
@patch.dict(os.environ, {"SHELL": "/bin/zsh", "PATH": "/fallback/path"})
|
|
1381
|
+
@patch("subprocess.run")
|
|
1382
|
+
def test_get_shell_path_unix_with_custom_shell(self, mock_run):
|
|
1383
|
+
"""Test get_shell_path() uses custom shell from environment."""
|
|
1384
|
+
mock_result = MagicMock()
|
|
1385
|
+
mock_result.returncode = 0
|
|
1386
|
+
mock_result.stdout = "/custom/shell/path\n"
|
|
1387
|
+
mock_run.return_value = mock_result
|
|
1388
|
+
|
|
1389
|
+
manager = MCPSessionManager()
|
|
1390
|
+
|
|
1391
|
+
result = manager.get_shell_path()
|
|
1392
|
+
|
|
1393
|
+
assert result == "/custom/shell/path"
|
|
1394
|
+
mock_run.assert_called_once_with(
|
|
1395
|
+
["/bin/zsh", "-l", "-c", "echo $PATH"],
|
|
1396
|
+
capture_output=True,
|
|
1397
|
+
text=True,
|
|
1398
|
+
timeout=3,
|
|
1399
|
+
)
|
|
1400
|
+
|
|
1401
|
+
@patch("sys.platform", "linux")
|
|
1402
|
+
@patch.dict(os.environ, {"PATH": "/fallback/path"}, clear=True)
|
|
1403
|
+
@patch("subprocess.run")
|
|
1404
|
+
def test_get_shell_path_unix_default_shell(self, mock_run):
|
|
1405
|
+
"""Test get_shell_path() uses default bash when SHELL not set."""
|
|
1406
|
+
mock_result = MagicMock()
|
|
1407
|
+
mock_result.returncode = 0
|
|
1408
|
+
mock_result.stdout = "/default/bash/path\n"
|
|
1409
|
+
mock_run.return_value = mock_result
|
|
1410
|
+
|
|
1411
|
+
manager = MCPSessionManager()
|
|
1412
|
+
|
|
1413
|
+
result = manager.get_shell_path()
|
|
1414
|
+
|
|
1415
|
+
assert result == "/default/bash/path"
|
|
1416
|
+
mock_run.assert_called_once_with(
|
|
1417
|
+
["/bin/bash", "-l", "-c", "echo $PATH"],
|
|
1418
|
+
capture_output=True,
|
|
1419
|
+
text=True,
|
|
1420
|
+
timeout=3,
|
|
1421
|
+
)
|
|
1422
|
+
|
|
1423
|
+
@patch("sys.platform", "linux")
|
|
1424
|
+
@patch.dict(os.environ, {"SHELL": "/bin/bash", "PATH": "/fallback/path"})
|
|
1425
|
+
@patch("subprocess.run")
|
|
1426
|
+
def test_get_shell_path_unix_subprocess_failure(self, mock_run):
|
|
1427
|
+
"""Test get_shell_path() falls back to environment PATH on subprocess failure."""
|
|
1428
|
+
# Mock failed subprocess execution
|
|
1429
|
+
mock_result = MagicMock()
|
|
1430
|
+
mock_result.returncode = 1
|
|
1431
|
+
mock_run.return_value = mock_result
|
|
1432
|
+
|
|
1433
|
+
manager = MCPSessionManager()
|
|
1434
|
+
|
|
1435
|
+
with patch("kiln_ai.tools.mcp_session_manager.logger") as mock_logger:
|
|
1436
|
+
result = manager.get_shell_path()
|
|
1437
|
+
|
|
1438
|
+
assert result == "/fallback/path"
|
|
1439
|
+
mock_logger.error.assert_called_once()
|
|
1440
|
+
assert "Error getting shell PATH" in mock_logger.error.call_args[0][0]
|
|
1441
|
+
|
|
1442
|
+
@patch("sys.platform", "linux")
|
|
1443
|
+
@patch.dict(os.environ, {"SHELL": "/bin/bash", "PATH": "/fallback/path"})
|
|
1444
|
+
@patch("subprocess.run")
|
|
1445
|
+
def test_get_shell_path_unix_subprocess_timeout(self, mock_run):
|
|
1446
|
+
"""Test get_shell_path() handles subprocess timeout."""
|
|
1447
|
+
# Mock subprocess timeout
|
|
1448
|
+
mock_run.side_effect = subprocess.TimeoutExpired(["bash"], 3)
|
|
1449
|
+
|
|
1450
|
+
manager = MCPSessionManager()
|
|
1451
|
+
|
|
1452
|
+
with patch("kiln_ai.tools.mcp_session_manager.logger") as mock_logger:
|
|
1453
|
+
result = manager.get_shell_path()
|
|
1454
|
+
|
|
1455
|
+
assert result == "/fallback/path"
|
|
1456
|
+
mock_logger.error.assert_any_call(
|
|
1457
|
+
"Shell path exception details: Command '['bash']' timed out after 3 seconds"
|
|
1458
|
+
)
|
|
1459
|
+
mock_logger.error.assert_any_call(
|
|
1460
|
+
"Error getting shell PATH. You may not be able to find MCP server commands like 'npx'. You can set a custom MCP path in the Kiln config file. See docs for details."
|
|
1461
|
+
)
|
|
1462
|
+
|
|
1463
|
+
@patch("sys.platform", "linux")
|
|
1464
|
+
@patch.dict(os.environ, {"SHELL": "/bin/bash", "PATH": "/fallback/path"})
|
|
1465
|
+
@patch("subprocess.run")
|
|
1466
|
+
def test_get_shell_path_unix_subprocess_error(self, mock_run):
|
|
1467
|
+
"""Test get_shell_path() handles subprocess errors."""
|
|
1468
|
+
# Mock subprocess error
|
|
1469
|
+
mock_run.side_effect = subprocess.SubprocessError("Command failed")
|
|
1470
|
+
|
|
1471
|
+
manager = MCPSessionManager()
|
|
1472
|
+
|
|
1473
|
+
with patch("kiln_ai.tools.mcp_session_manager.logger") as mock_logger:
|
|
1474
|
+
result = manager.get_shell_path()
|
|
1475
|
+
|
|
1476
|
+
assert result == "/fallback/path"
|
|
1477
|
+
mock_logger.error.assert_any_call(
|
|
1478
|
+
"Shell path exception details: Command failed"
|
|
1479
|
+
)
|
|
1480
|
+
|
|
1481
|
+
@patch("sys.platform", "linux")
|
|
1482
|
+
@patch.dict(os.environ, {"SHELL": "/bin/bash", "PATH": "/fallback/path"})
|
|
1483
|
+
@patch("subprocess.run")
|
|
1484
|
+
def test_get_shell_path_unix_general_exception(self, mock_run):
|
|
1485
|
+
"""Test get_shell_path() handles general exceptions."""
|
|
1486
|
+
# Mock general exception
|
|
1487
|
+
mock_run.side_effect = RuntimeError("Unexpected error")
|
|
1488
|
+
|
|
1489
|
+
manager = MCPSessionManager()
|
|
1490
|
+
|
|
1491
|
+
with patch("kiln_ai.tools.mcp_session_manager.logger") as mock_logger:
|
|
1492
|
+
result = manager.get_shell_path()
|
|
1493
|
+
|
|
1494
|
+
assert result == "/fallback/path"
|
|
1495
|
+
mock_logger.error.assert_any_call(
|
|
1496
|
+
"Shell path exception details: Unexpected error"
|
|
1497
|
+
)
|
|
1498
|
+
|
|
1499
|
+
@patch("sys.platform", "linux")
|
|
1500
|
+
@patch.dict(os.environ, {"SHELL": "/bin/bash"}, clear=True)
|
|
1501
|
+
@patch("subprocess.run")
|
|
1502
|
+
def test_get_shell_path_unix_no_fallback_path(self, mock_run):
|
|
1503
|
+
"""Test get_shell_path() when no PATH environment variable exists."""
|
|
1504
|
+
mock_run.side_effect = subprocess.SubprocessError("Command failed")
|
|
1505
|
+
|
|
1506
|
+
manager = MCPSessionManager()
|
|
1507
|
+
|
|
1508
|
+
result = manager.get_shell_path()
|
|
1509
|
+
|
|
1510
|
+
assert result == ""
|
|
1511
|
+
|
|
1512
|
+
@patch("sys.platform", "linux")
|
|
1513
|
+
@patch.dict(os.environ, {"SHELL": "/bin/bash", "PATH": "/original/path"})
|
|
1514
|
+
@patch("subprocess.run")
|
|
1515
|
+
def test_get_shell_path_caching(self, mock_run):
|
|
1516
|
+
"""Test get_shell_path() caches the result."""
|
|
1517
|
+
mock_result = MagicMock()
|
|
1518
|
+
mock_result.returncode = 0
|
|
1519
|
+
mock_result.stdout = "/cached/path\n"
|
|
1520
|
+
mock_run.return_value = mock_result
|
|
1521
|
+
|
|
1522
|
+
manager = MCPSessionManager()
|
|
1523
|
+
|
|
1524
|
+
# First call should execute subprocess
|
|
1525
|
+
result1 = manager.get_shell_path()
|
|
1526
|
+
assert result1 == "/cached/path"
|
|
1527
|
+
assert mock_run.call_count == 1
|
|
1528
|
+
|
|
1529
|
+
# Second call should use cached value
|
|
1530
|
+
result2 = manager.get_shell_path()
|
|
1531
|
+
assert result2 == "/cached/path"
|
|
1532
|
+
assert mock_run.call_count == 1 # Should not have been called again
|
|
1533
|
+
|
|
1534
|
+
|
|
1535
|
+
class TestMCPServerIntegration:
|
|
1536
|
+
"""Integration tests for MCPServer using real services."""
|
|
1537
|
+
|
|
1538
|
+
@pytest.mark.skip(
|
|
1539
|
+
reason="Skipping integration test since it requires calling a real MCP server"
|
|
1540
|
+
)
|
|
1541
|
+
async def test_list_tools_with_real_remote_mcp_server(self):
|
|
1542
|
+
"""Test list_tools with a real MCP server if available."""
|
|
1543
|
+
external_tool_server = ExternalToolServer(
|
|
1544
|
+
name="postman_echo",
|
|
1545
|
+
type=ToolServerType.remote_mcp,
|
|
1546
|
+
description="Postman Echo MCP Server for testing",
|
|
1547
|
+
properties={
|
|
1548
|
+
"server_url": "https://postman-echo-mcp.fly.dev/",
|
|
1549
|
+
"headers": {},
|
|
1550
|
+
},
|
|
1551
|
+
)
|
|
1552
|
+
|
|
1553
|
+
async with MCPSessionManager.shared().mcp_client(
|
|
1554
|
+
external_tool_server
|
|
1555
|
+
) as session:
|
|
1556
|
+
tools = await session.list_tools()
|
|
1557
|
+
|
|
1558
|
+
assert tools is not None
|
|
1559
|
+
assert len(tools.tools) > 0
|
|
1560
|
+
assert "echo" in [tool.name for tool in tools.tools]
|
|
1561
|
+
|
|
1562
|
+
@pytest.mark.skip(
|
|
1563
|
+
reason="Skipping integration test since it requires calling a real MCP server"
|
|
1564
|
+
)
|
|
1565
|
+
async def test_list_tools_with_real_local_mcp_server(self):
|
|
1566
|
+
"""Test list_tools with a real local MCP server if available."""
|
|
1567
|
+
external_tool_server = ExternalToolServer(
|
|
1568
|
+
name="Firecrawl",
|
|
1569
|
+
type=ToolServerType.local_mcp,
|
|
1570
|
+
description="Firecrawl MCP Server for testing",
|
|
1571
|
+
properties={
|
|
1572
|
+
"command": "npx",
|
|
1573
|
+
"args": ["-y", "firecrawl-mcp"],
|
|
1574
|
+
"env_vars": {"FIRECRAWL_API_KEY": "REPLACE_WITH_YOUR_API_KEY"},
|
|
1575
|
+
},
|
|
1576
|
+
)
|
|
1577
|
+
|
|
1578
|
+
async with MCPSessionManager.shared().mcp_client(
|
|
1579
|
+
external_tool_server
|
|
1580
|
+
) as session:
|
|
1581
|
+
tools = await session.list_tools()
|
|
1582
|
+
|
|
1583
|
+
assert tools is not None
|
|
1584
|
+
assert len(tools.tools) > 0
|
|
1585
|
+
assert "firecrawl_scrape" in [tool.name for tool in tools.tools]
|