mcpower-proxy 0.0.65__py3-none-any.whl → 0.0.74__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 mcpower-proxy might be problematic. Click here for more details.
- ide_tools/__init__.py +12 -0
- ide_tools/common/__init__.py +5 -0
- ide_tools/common/hooks/__init__.py +5 -0
- ide_tools/common/hooks/init.py +124 -0
- ide_tools/common/hooks/output.py +63 -0
- ide_tools/common/hooks/prompt_submit.py +133 -0
- ide_tools/common/hooks/read_file.py +167 -0
- ide_tools/common/hooks/shell_execution.py +255 -0
- ide_tools/common/hooks/shell_parser_bashlex.py +277 -0
- ide_tools/common/hooks/types.py +34 -0
- ide_tools/common/hooks/utils.py +286 -0
- ide_tools/cursor/__init__.py +11 -0
- ide_tools/cursor/constants.py +58 -0
- ide_tools/cursor/format.py +35 -0
- ide_tools/cursor/router.py +100 -0
- ide_tools/router.py +48 -0
- main.py +11 -4
- {mcpower_proxy-0.0.65.dist-info → mcpower_proxy-0.0.74.dist-info}/METADATA +4 -3
- mcpower_proxy-0.0.74.dist-info/RECORD +60 -0
- {mcpower_proxy-0.0.65.dist-info → mcpower_proxy-0.0.74.dist-info}/top_level.txt +1 -0
- modules/apis/security_policy.py +11 -6
- modules/decision_handler.py +219 -0
- modules/logs/audit_trail.py +16 -15
- modules/logs/logger.py +14 -18
- modules/redaction/gitleaks_rules.py +1 -1
- modules/redaction/pii_rules.py +0 -48
- modules/redaction/redactor.py +112 -107
- modules/ui/__init__.py +1 -1
- modules/ui/confirmation.py +0 -1
- modules/utils/cli.py +36 -6
- modules/utils/ids.py +55 -10
- modules/utils/json.py +3 -3
- wrapper/__version__.py +1 -1
- wrapper/middleware.py +135 -217
- wrapper/server.py +19 -11
- mcpower_proxy-0.0.65.dist-info/RECORD +0 -43
- {mcpower_proxy-0.0.65.dist-info → mcpower_proxy-0.0.74.dist-info}/WHEEL +0 -0
- {mcpower_proxy-0.0.65.dist-info → mcpower_proxy-0.0.74.dist-info}/entry_points.txt +0 -0
- {mcpower_proxy-0.0.65.dist-info → mcpower_proxy-0.0.74.dist-info}/licenses/LICENSE +0 -0
modules/utils/ids.py
CHANGED
|
@@ -2,6 +2,7 @@
|
|
|
2
2
|
Utilities for generating event IDs, session IDs, app UIDs, and timing helpers
|
|
3
3
|
"""
|
|
4
4
|
import os
|
|
5
|
+
import sys
|
|
5
6
|
import time
|
|
6
7
|
import uuid
|
|
7
8
|
from pathlib import Path
|
|
@@ -23,6 +24,16 @@ def generate_event_id() -> str:
|
|
|
23
24
|
return f"{timestamp}-{unique_part}"
|
|
24
25
|
|
|
25
26
|
|
|
27
|
+
def generate_prompt_id() -> str:
|
|
28
|
+
"""
|
|
29
|
+
Generate truly-random 8-character prompt ID for user request correlation
|
|
30
|
+
|
|
31
|
+
Returns:
|
|
32
|
+
8-character random ID string
|
|
33
|
+
"""
|
|
34
|
+
return str(uuid.uuid4())[:8]
|
|
35
|
+
|
|
36
|
+
|
|
26
37
|
def get_session_id() -> str:
|
|
27
38
|
"""
|
|
28
39
|
Get session ID for the current process. Returns the same value for all calls
|
|
@@ -67,7 +78,8 @@ def _atomic_write_uuid(file_path: Path, new_uuid: str) -> bool:
|
|
|
67
78
|
True if write succeeded, False if file exists
|
|
68
79
|
"""
|
|
69
80
|
try:
|
|
70
|
-
|
|
81
|
+
mode = 0o666 if sys.platform == 'win32' else 0o600
|
|
82
|
+
fd = os.open(str(file_path), os.O_CREAT | os.O_EXCL | os.O_WRONLY, mode)
|
|
71
83
|
try:
|
|
72
84
|
os.write(fd, new_uuid.encode('utf-8'))
|
|
73
85
|
finally:
|
|
@@ -91,7 +103,7 @@ def _get_or_create_uuid(uid_path: Path, logger, id_type: str) -> str:
|
|
|
91
103
|
UUID string
|
|
92
104
|
"""
|
|
93
105
|
uid_path.parent.mkdir(parents=True, exist_ok=True)
|
|
94
|
-
|
|
106
|
+
|
|
95
107
|
max_attempts = 3
|
|
96
108
|
for attempt in range(max_attempts):
|
|
97
109
|
if uid_path.exists():
|
|
@@ -107,20 +119,53 @@ def _get_or_create_uuid(uid_path: Path, logger, id_type: str) -> str:
|
|
|
107
119
|
time.sleep(0.1 * (2 ** attempt))
|
|
108
120
|
continue
|
|
109
121
|
raise
|
|
110
|
-
|
|
122
|
+
|
|
111
123
|
new_uid = str(uuid.uuid4())
|
|
112
|
-
|
|
124
|
+
|
|
113
125
|
if _atomic_write_uuid(uid_path, new_uid):
|
|
114
|
-
logger.info(f"Generated {id_type}: {new_uid}")
|
|
126
|
+
logger.info(f"Generated {id_type}: {new_uid} at {uid_path}")
|
|
115
127
|
return new_uid
|
|
116
|
-
|
|
117
|
-
logger.debug(
|
|
128
|
+
|
|
129
|
+
logger.debug(
|
|
130
|
+
f"{id_type.title()} file created by another process, reading (attempt {attempt + 1}/{max_attempts})")
|
|
118
131
|
if attempt < max_attempts - 1:
|
|
119
132
|
time.sleep(0.05)
|
|
120
|
-
|
|
133
|
+
|
|
121
134
|
raise RuntimeError(f"Failed to get or create {id_type} after {max_attempts} attempts")
|
|
122
135
|
|
|
123
136
|
|
|
137
|
+
def get_home_mcpower_dir() -> Path:
|
|
138
|
+
"""
|
|
139
|
+
Get the global MCPower directory path in user's home directory
|
|
140
|
+
|
|
141
|
+
Returns:
|
|
142
|
+
Path to ~/.mcpower directory
|
|
143
|
+
"""
|
|
144
|
+
return Path.home() / ".mcpower"
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def get_project_mcpower_dir(project_path: Optional[str] = None) -> str:
|
|
148
|
+
"""
|
|
149
|
+
Get the MCPower directory path, with fallback to global ~/.mcpower
|
|
150
|
+
|
|
151
|
+
Args:
|
|
152
|
+
project_path: Optional project/workspace path. If None or invalid, falls back to ~/.mcpower
|
|
153
|
+
|
|
154
|
+
Returns:
|
|
155
|
+
Path to use for MCPower data (either project/.mcpower or ~/.mcpower)
|
|
156
|
+
"""
|
|
157
|
+
if project_path:
|
|
158
|
+
try:
|
|
159
|
+
path = Path(project_path)
|
|
160
|
+
if path.exists() and path.is_dir():
|
|
161
|
+
return str(path)
|
|
162
|
+
except Exception:
|
|
163
|
+
pass
|
|
164
|
+
|
|
165
|
+
# Fallback to global ~/.mcpower
|
|
166
|
+
return str(get_home_mcpower_dir())
|
|
167
|
+
|
|
168
|
+
|
|
124
169
|
def get_or_create_user_id(logger) -> str:
|
|
125
170
|
"""
|
|
126
171
|
Get or create machine-wide user ID from ~/.mcpower/uid
|
|
@@ -132,7 +177,7 @@ def get_or_create_user_id(logger) -> str:
|
|
|
132
177
|
Returns:
|
|
133
178
|
User ID string
|
|
134
179
|
"""
|
|
135
|
-
uid_path =
|
|
180
|
+
uid_path = get_home_mcpower_dir() / "uid"
|
|
136
181
|
return _get_or_create_uuid(uid_path, logger, "user ID")
|
|
137
182
|
|
|
138
183
|
|
|
@@ -156,5 +201,5 @@ def read_app_uid(logger, project_folder_path: str) -> str:
|
|
|
156
201
|
else:
|
|
157
202
|
# Project-specific case
|
|
158
203
|
uid_path = project_path / ".mcpower" / "app_uid"
|
|
159
|
-
|
|
204
|
+
|
|
160
205
|
return _get_or_create_uuid(uid_path, logger, "app UID")
|
modules/utils/json.py
CHANGED
|
@@ -52,7 +52,7 @@ def safe_json_dumps(obj: Any, **kwargs) -> str:
|
|
|
52
52
|
# If it's a Pydantic BaseModel, use its built-in JSON serialization
|
|
53
53
|
if isinstance(obj, BaseModel):
|
|
54
54
|
return obj.model_dump_json(**kwargs)
|
|
55
|
-
|
|
55
|
+
|
|
56
56
|
# If it's a dict or list that might contain Pydantic objects, use custom serializer
|
|
57
57
|
def default_serializer(o):
|
|
58
58
|
if isinstance(o, BaseModel):
|
|
@@ -72,7 +72,7 @@ def safe_json_dumps(obj: Any, **kwargs) -> str:
|
|
|
72
72
|
return o.__dict__
|
|
73
73
|
# Fallback to string representation
|
|
74
74
|
return str(o)
|
|
75
|
-
|
|
75
|
+
|
|
76
76
|
return json.dumps(obj, default=default_serializer, **kwargs)
|
|
77
77
|
|
|
78
78
|
|
|
@@ -117,4 +117,4 @@ def parse_jsonc(text: str) -> Any:
|
|
|
117
117
|
return json.loads(text)
|
|
118
118
|
except json.JSONDecodeError:
|
|
119
119
|
# Re-raise the original JSONC error if JSON also fails
|
|
120
|
-
raise json.JSONDecodeError(f"JSONC parsing failed: {str(e)}", text, 0)
|
|
120
|
+
raise json.JSONDecodeError(f"JSONC parsing failed: {str(e)}", text, 0)
|
wrapper/__version__.py
CHANGED
wrapper/middleware.py
CHANGED
|
@@ -2,6 +2,8 @@
|
|
|
2
2
|
FastMCP middleware for security policy enforcement
|
|
3
3
|
Implements pre/post interception for all MCP operations
|
|
4
4
|
"""
|
|
5
|
+
import asyncio
|
|
6
|
+
import sys
|
|
5
7
|
import time
|
|
6
8
|
import urllib.parse
|
|
7
9
|
from datetime import datetime, timezone
|
|
@@ -11,22 +13,23 @@ from typing import Any, Dict, List, Optional
|
|
|
11
13
|
from fastmcp.exceptions import FastMCPError
|
|
12
14
|
from fastmcp.server.middleware.middleware import Middleware, MiddlewareContext, CallNext
|
|
13
15
|
from fastmcp.server.proxy import ProxyClient
|
|
16
|
+
from httpx import HTTPStatusError
|
|
17
|
+
from mcp import ErrorData
|
|
18
|
+
|
|
19
|
+
from mcpower_shared.mcp_types import (create_policy_request, create_policy_response, AgentContext, EnvironmentContext,
|
|
20
|
+
InitRequest,
|
|
21
|
+
ServerRef, ToolRef)
|
|
14
22
|
from modules.apis.security_policy import SecurityPolicyClient
|
|
23
|
+
from modules.decision_handler import DecisionHandler, DecisionEnforcementError
|
|
15
24
|
from modules.logs.audit_trail import AuditTrailLogger
|
|
16
25
|
from modules.logs.logger import MCPLogger
|
|
17
26
|
from modules.redaction import redact
|
|
18
|
-
from modules.ui.classes import ConfirmationRequest, DialogOptions, UserDecision
|
|
19
|
-
from modules.ui.confirmation import UserConfirmationDialog, UserConfirmationError
|
|
20
27
|
from modules.utils.copy import safe_copy
|
|
21
|
-
from modules.utils.ids import generate_event_id, get_session_id, read_app_uid
|
|
28
|
+
from modules.utils.ids import generate_event_id, get_session_id, read_app_uid, get_project_mcpower_dir
|
|
22
29
|
from modules.utils.json import safe_json_dumps, to_dict
|
|
23
30
|
from modules.utils.mcp_configs import extract_wrapped_server_info
|
|
24
31
|
from wrapper.schema import merge_input_schema_with_existing
|
|
25
32
|
|
|
26
|
-
from mcpower_shared.mcp_types import (create_policy_request, create_policy_response, AgentContext, EnvironmentContext,
|
|
27
|
-
InitRequest,
|
|
28
|
-
ServerRef, ToolRef, UserConfirmation)
|
|
29
|
-
|
|
30
33
|
|
|
31
34
|
class MockContext:
|
|
32
35
|
"""Mock context for internal operations"""
|
|
@@ -52,10 +55,7 @@ class MockContext:
|
|
|
52
55
|
class SecurityMiddleware(Middleware):
|
|
53
56
|
"""FastMCP middleware for security policy enforcement"""
|
|
54
57
|
|
|
55
|
-
app_id: str = ""
|
|
56
58
|
_TOOLS_INIT_DEBOUNCE_SECONDS = 60
|
|
57
|
-
_last_tools_init_time: Optional[float] = None
|
|
58
|
-
_last_workspace_root: Optional[str] = None
|
|
59
59
|
|
|
60
60
|
def __init__(self,
|
|
61
61
|
wrapped_server_configs: dict,
|
|
@@ -71,6 +71,9 @@ class SecurityMiddleware(Middleware):
|
|
|
71
71
|
self.audit_logger = audit_logger
|
|
72
72
|
self.app_id = ""
|
|
73
73
|
self._last_workspace_root = None
|
|
74
|
+
self._last_tools_init_time: Optional[float] = None
|
|
75
|
+
self._tools_list_in_progress: Optional[asyncio.Task] = None
|
|
76
|
+
self._tools_list_lock = asyncio.Lock()
|
|
74
77
|
|
|
75
78
|
self.wrapped_server_name, self.wrapped_server_transport = (
|
|
76
79
|
extract_wrapped_server_info(self.wrapper_server_name, self.logger, self.wrapped_server_configs)
|
|
@@ -87,17 +90,36 @@ class SecurityMiddleware(Middleware):
|
|
|
87
90
|
async def on_message(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
|
88
91
|
self.logger.info(f"on_message: {redact(safe_json_dumps(context))}")
|
|
89
92
|
|
|
90
|
-
#
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
if
|
|
94
|
-
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
self.
|
|
93
|
+
# Skip workspace check for `initialize` calls to avoid premature app_uid changes.
|
|
94
|
+
# The `initialize` request doesn't contain workspace data, so checking it would
|
|
95
|
+
# cause unnecessary audit log flushes before the actual workspace init arrives.
|
|
96
|
+
if context.method != "initialize":
|
|
97
|
+
# Check workspace roots and re-initialize app_uid if workspace changed
|
|
98
|
+
workspace_roots = await self._extract_workspace_roots(context)
|
|
99
|
+
current_workspace_root = get_project_mcpower_dir(workspace_roots[0] if workspace_roots else None)
|
|
100
|
+
if current_workspace_root != self._last_workspace_root:
|
|
101
|
+
self.logger.debug(
|
|
102
|
+
f"Workspace root changed from {self._last_workspace_root} to {current_workspace_root}")
|
|
103
|
+
self._last_workspace_root = current_workspace_root
|
|
104
|
+
self.app_id = read_app_uid(logger=self.logger, project_folder_path=current_workspace_root)
|
|
105
|
+
self.audit_logger.set_app_uid(self.app_id)
|
|
98
106
|
|
|
99
107
|
operation_type = "message"
|
|
100
|
-
|
|
108
|
+
|
|
109
|
+
async def call_next_wrapper(ctx):
|
|
110
|
+
try:
|
|
111
|
+
return await call_next(ctx)
|
|
112
|
+
except HTTPStatusError as e:
|
|
113
|
+
if e.response.status_code in (401, 403):
|
|
114
|
+
raise FastMCPError(ErrorData(
|
|
115
|
+
code=-32000,
|
|
116
|
+
message="Authentication required",
|
|
117
|
+
data={
|
|
118
|
+
"type": "unauthorized",
|
|
119
|
+
"details": "Please provide valid authentication credentials"
|
|
120
|
+
}
|
|
121
|
+
))
|
|
122
|
+
raise e
|
|
101
123
|
|
|
102
124
|
match context.type:
|
|
103
125
|
case "request":
|
|
@@ -114,13 +136,13 @@ class SecurityMiddleware(Middleware):
|
|
|
114
136
|
operation_type = "prompt"
|
|
115
137
|
case "tools/list":
|
|
116
138
|
# Special handling for tools/list - call /init instead of normal inspection
|
|
117
|
-
return await self._handle_tools_list(context,
|
|
118
|
-
case "resources/list" | "resources/templates/list" | "prompts/list":
|
|
119
|
-
return await
|
|
139
|
+
return await self._handle_tools_list(context, call_next_wrapper)
|
|
140
|
+
case "initialize" | "resources/list" | "resources/templates/list" | "prompts/list":
|
|
141
|
+
return await call_next_wrapper(context)
|
|
120
142
|
|
|
121
143
|
return await self._handle_operation(
|
|
122
144
|
context=context,
|
|
123
|
-
call_next=
|
|
145
|
+
call_next=call_next_wrapper,
|
|
124
146
|
error_class=FastMCPError,
|
|
125
147
|
operation_type=operation_type
|
|
126
148
|
)
|
|
@@ -180,15 +202,15 @@ class SecurityMiddleware(Middleware):
|
|
|
180
202
|
return await ProxyClient.default_progress_handler(progress, total, message)
|
|
181
203
|
|
|
182
204
|
async def secure_log_handler(self, log_message):
|
|
183
|
-
# FIXME: log_message should be redacted before logging,
|
|
205
|
+
# FIXME: log_message should be redacted before logging,
|
|
184
206
|
self.logger.info(f"secure_log_handler: {str(log_message)[:100]}...")
|
|
185
207
|
# FIXME: log_message should be reviewed with policy before forwarding
|
|
186
|
-
|
|
208
|
+
|
|
187
209
|
# Handle case where log_message.data is a string instead of dict
|
|
188
210
|
# The default_log_handler expects data to be a dict with 'msg' and 'extra' keys
|
|
189
211
|
if hasattr(log_message, 'data') and isinstance(log_message.data, str):
|
|
190
212
|
log_message = safe_copy(log_message, {'data': {'msg': log_message.data, 'extra': None}})
|
|
191
|
-
|
|
213
|
+
|
|
192
214
|
return await ProxyClient.default_log_handler(log_message)
|
|
193
215
|
|
|
194
216
|
async def _handle_operation(self, context: MiddlewareContext, call_next, error_class, operation_type: str):
|
|
@@ -221,19 +243,28 @@ class SecurityMiddleware(Middleware):
|
|
|
221
243
|
prompt_id=prompt_id
|
|
222
244
|
)
|
|
223
245
|
on_inspect_request_duration = time.time() - on_inspect_request_start_time
|
|
224
|
-
self.logger.
|
|
246
|
+
self.logger.debug(
|
|
247
|
+
f"PROFILE: {operation_type} id: {event_id} inspect_request duration: {on_inspect_request_duration:.2f} seconds")
|
|
225
248
|
|
|
226
|
-
|
|
227
|
-
|
|
228
|
-
|
|
229
|
-
|
|
230
|
-
|
|
231
|
-
|
|
232
|
-
|
|
233
|
-
|
|
234
|
-
|
|
235
|
-
|
|
236
|
-
|
|
249
|
+
try:
|
|
250
|
+
await DecisionHandler(
|
|
251
|
+
logger=self.logger,
|
|
252
|
+
audit_logger=self.audit_logger,
|
|
253
|
+
session_id=self.session_id,
|
|
254
|
+
app_id=self.app_id
|
|
255
|
+
).enforce_decision(
|
|
256
|
+
decision=request_decision,
|
|
257
|
+
is_request=True,
|
|
258
|
+
event_id=event_id,
|
|
259
|
+
tool_name=tool_name,
|
|
260
|
+
content_data=tool_args,
|
|
261
|
+
operation_type=operation_type,
|
|
262
|
+
prompt_id=prompt_id,
|
|
263
|
+
server_name=self.wrapped_server_name,
|
|
264
|
+
error_message_prefix=f"{operation_type.title()} request blocked by security policy"
|
|
265
|
+
)
|
|
266
|
+
except DecisionEnforcementError as e:
|
|
267
|
+
raise error_class(str(e))
|
|
237
268
|
|
|
238
269
|
self.audit_logger.log_event(
|
|
239
270
|
"agent_request_forwarded",
|
|
@@ -250,7 +281,8 @@ class SecurityMiddleware(Middleware):
|
|
|
250
281
|
# Call wrapped MCP with cleaned context (e.g., no wrapper args)
|
|
251
282
|
result = await call_next(cleaned_context)
|
|
252
283
|
on_call_next_duration = time.time() - on_call_next_start_time
|
|
253
|
-
self.logger.
|
|
284
|
+
self.logger.debug(
|
|
285
|
+
f"PROFILE: {operation_type} id: {event_id} call_next duration: {on_call_next_duration:.2f} seconds")
|
|
254
286
|
|
|
255
287
|
response_content = self._extract_response_content(result)
|
|
256
288
|
|
|
@@ -275,19 +307,28 @@ class SecurityMiddleware(Middleware):
|
|
|
275
307
|
prompt_id=prompt_id
|
|
276
308
|
)
|
|
277
309
|
on_inspect_response_duration = time.time() - on_inspect_response_start_time
|
|
278
|
-
self.logger.
|
|
310
|
+
self.logger.debug(
|
|
311
|
+
f"PROFILE: {operation_type} id: {event_id} inspect_response duration: {on_inspect_response_duration:.2f} seconds")
|
|
279
312
|
|
|
280
|
-
|
|
281
|
-
|
|
282
|
-
|
|
283
|
-
|
|
284
|
-
|
|
285
|
-
|
|
286
|
-
|
|
287
|
-
|
|
288
|
-
|
|
289
|
-
|
|
290
|
-
|
|
313
|
+
try:
|
|
314
|
+
await DecisionHandler(
|
|
315
|
+
logger=self.logger,
|
|
316
|
+
audit_logger=self.audit_logger,
|
|
317
|
+
session_id=self.session_id,
|
|
318
|
+
app_id=self.app_id
|
|
319
|
+
).enforce_decision(
|
|
320
|
+
decision=response_decision,
|
|
321
|
+
is_request=False,
|
|
322
|
+
event_id=event_id,
|
|
323
|
+
tool_name=tool_name,
|
|
324
|
+
content_data=response_content,
|
|
325
|
+
operation_type=operation_type,
|
|
326
|
+
prompt_id=prompt_id,
|
|
327
|
+
server_name=self.wrapped_server_name,
|
|
328
|
+
error_message_prefix=f"{operation_type.title()} response blocked by security policy"
|
|
329
|
+
)
|
|
330
|
+
except DecisionEnforcementError as e:
|
|
331
|
+
raise error_class(str(e))
|
|
291
332
|
|
|
292
333
|
self.audit_logger.log_event(
|
|
293
334
|
"mcp_response_forwarded",
|
|
@@ -300,15 +341,30 @@ class SecurityMiddleware(Middleware):
|
|
|
300
341
|
prompt_id=prompt_id
|
|
301
342
|
)
|
|
302
343
|
on_handle_operation_duration = time.time() - on_handle_operation_start_time
|
|
303
|
-
self.logger.
|
|
344
|
+
self.logger.debug(
|
|
345
|
+
f"PROFILE: {operation_type} id: {event_id} duration: {on_handle_operation_duration:.2f} seconds")
|
|
304
346
|
return result
|
|
305
347
|
|
|
306
348
|
async def _handle_tools_list(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
|
307
|
-
"""Handle tools/list by calling /init API and modifying schemas"""
|
|
349
|
+
"""Handle tools/list by calling /init API and modifying schemas with deduplication"""
|
|
308
350
|
event_id = generate_event_id()
|
|
309
351
|
on_handle_tools_list_start_time = time.time()
|
|
310
|
-
|
|
311
|
-
|
|
352
|
+
|
|
353
|
+
async with self._tools_list_lock:
|
|
354
|
+
if not self._tools_list_in_progress or self._tools_list_in_progress.done():
|
|
355
|
+
self._tools_list_in_progress = asyncio.create_task(call_next(context))
|
|
356
|
+
shared_task = self._tools_list_in_progress
|
|
357
|
+
|
|
358
|
+
try:
|
|
359
|
+
result = await shared_task
|
|
360
|
+
except Exception as e:
|
|
361
|
+
async with self._tools_list_lock:
|
|
362
|
+
if self._tools_list_in_progress is shared_task:
|
|
363
|
+
self._tools_list_in_progress = None
|
|
364
|
+
raise
|
|
365
|
+
self.logger.debug(
|
|
366
|
+
f"PROFILE: tools/list call_next duration: {time.time() - on_handle_tools_list_start_time:.2f} seconds id: {event_id}")
|
|
367
|
+
|
|
312
368
|
tools_list = None
|
|
313
369
|
if isinstance(result, list):
|
|
314
370
|
tools_list = result
|
|
@@ -338,11 +394,13 @@ class SecurityMiddleware(Middleware):
|
|
|
338
394
|
enhanced_result = result
|
|
339
395
|
|
|
340
396
|
on_handle_tools_list_duration = time.time() - on_handle_tools_list_start_time
|
|
341
|
-
self.logger.
|
|
397
|
+
self.logger.debug(
|
|
398
|
+
f"PROFILE: tools/list enhanced_result duration: {on_handle_tools_list_duration:.2f} seconds id: {event_id}")
|
|
342
399
|
return enhanced_result
|
|
343
400
|
|
|
344
401
|
on_handle_tools_list_duration = time.time() - on_handle_tools_list_start_time
|
|
345
|
-
self.logger.
|
|
402
|
+
self.logger.debug(
|
|
403
|
+
f"PROFILE: tools/list result duration: {on_handle_tools_list_duration:.2f} seconds id: {event_id}")
|
|
346
404
|
|
|
347
405
|
return result
|
|
348
406
|
|
|
@@ -481,6 +539,12 @@ class SecurityMiddleware(Middleware):
|
|
|
481
539
|
file_path_prefix = 'file://'
|
|
482
540
|
if uri.startswith(file_path_prefix):
|
|
483
541
|
path = urllib.parse.unquote(uri[len(file_path_prefix):])
|
|
542
|
+
|
|
543
|
+
# Windows fix: remove leading slash before drive letter
|
|
544
|
+
# file:///C:/path becomes /C:/path, should be C:/path
|
|
545
|
+
if sys.platform == 'win32' and len(path) >= 3 and path[0] == '/' and path[2] == ':':
|
|
546
|
+
path = path[1:]
|
|
547
|
+
|
|
484
548
|
try:
|
|
485
549
|
resolved_path = str(Path(path).resolve())
|
|
486
550
|
workspace_roots.append(resolved_path)
|
|
@@ -504,9 +568,13 @@ class SecurityMiddleware(Middleware):
|
|
|
504
568
|
base_dict = await self._build_baseline_policy_dict(event_id, context, wrapper_args, tool_args)
|
|
505
569
|
policy_request = create_policy_request(
|
|
506
570
|
event_id=event_id,
|
|
507
|
-
|
|
508
|
-
|
|
509
|
-
|
|
571
|
+
server=ServerRef(
|
|
572
|
+
name=base_dict["server"]["name"],
|
|
573
|
+
transport=base_dict["server"]["transport"]
|
|
574
|
+
),
|
|
575
|
+
tool=ToolRef(
|
|
576
|
+
name=base_dict["tool"]["name"] or base_dict["tool"]["method"]
|
|
577
|
+
),
|
|
510
578
|
agent_context=base_dict["agent_context"],
|
|
511
579
|
env_context=base_dict["environment_context"],
|
|
512
580
|
arguments=tool_args,
|
|
@@ -532,9 +600,13 @@ class SecurityMiddleware(Middleware):
|
|
|
532
600
|
base_dict = await self._build_baseline_policy_dict(event_id, context, wrapper_args, tool_args)
|
|
533
601
|
policy_response = create_policy_response(
|
|
534
602
|
event_id=event_id,
|
|
535
|
-
|
|
536
|
-
|
|
537
|
-
|
|
603
|
+
server=ServerRef(
|
|
604
|
+
name=base_dict["server"]["name"],
|
|
605
|
+
transport=base_dict["server"]["transport"]
|
|
606
|
+
),
|
|
607
|
+
tool=ToolRef(
|
|
608
|
+
name=base_dict["tool"]["name"] or base_dict["tool"]["method"]
|
|
609
|
+
),
|
|
538
610
|
response_content=safe_json_dumps(result),
|
|
539
611
|
agent_context=base_dict["agent_context"],
|
|
540
612
|
env_context=base_dict["environment_context"],
|
|
@@ -586,28 +658,6 @@ class SecurityMiddleware(Middleware):
|
|
|
586
658
|
)
|
|
587
659
|
}
|
|
588
660
|
|
|
589
|
-
async def _record_user_confirmation(self, event_id: str, is_request: bool, user_decision: UserDecision,
|
|
590
|
-
prompt_id: str, call_type: str = None):
|
|
591
|
-
"""Record user confirmation decision with the security API"""
|
|
592
|
-
try:
|
|
593
|
-
direction = "request" if is_request else "response"
|
|
594
|
-
|
|
595
|
-
user_confirmation = UserConfirmation(
|
|
596
|
-
event_id=event_id,
|
|
597
|
-
direction=direction,
|
|
598
|
-
user_decision=user_decision,
|
|
599
|
-
call_type=call_type
|
|
600
|
-
)
|
|
601
|
-
|
|
602
|
-
async with SecurityPolicyClient(session_id=self.session_id, logger=self.logger,
|
|
603
|
-
audit_logger=self.audit_logger, app_id=self.app_id) as client:
|
|
604
|
-
result = await client.record_user_confirmation(user_confirmation, prompt_id=prompt_id)
|
|
605
|
-
self.logger.debug(f"User confirmation recorded: {result}")
|
|
606
|
-
except Exception as e:
|
|
607
|
-
# Don't fail the operation if API call fails - just log the error
|
|
608
|
-
self.logger.error(f"Failed to record user confirmation: {e}")
|
|
609
|
-
|
|
610
|
-
|
|
611
661
|
@staticmethod
|
|
612
662
|
def _create_security_api_failure_decision(error: Exception) -> Dict[str, Any]:
|
|
613
663
|
"""Create a standard failure decision when security API is unavailable/failing/unreachable"""
|
|
@@ -617,135 +667,3 @@ class SecurityMiddleware(Middleware):
|
|
|
617
667
|
"reasons": [f"Security API unavailable: {error}"],
|
|
618
668
|
"matched_rules": ["security_api.error"]
|
|
619
669
|
}
|
|
620
|
-
|
|
621
|
-
async def _enforce_decision(self, decision: Dict[str, Any], error_class, base_message: str,
|
|
622
|
-
is_request: bool, event_id: str, tool_name: str, content_data: Dict[str, Any],
|
|
623
|
-
operation_type: str, prompt_id: str):
|
|
624
|
-
"""Enforce security decision with user confirmation support"""
|
|
625
|
-
decision_type = decision.get("decision", "block")
|
|
626
|
-
|
|
627
|
-
if decision_type == "allow":
|
|
628
|
-
return
|
|
629
|
-
|
|
630
|
-
elif decision_type == "block":
|
|
631
|
-
policy_reasons = decision.get("reasons", ["Policy violation"])
|
|
632
|
-
severity = decision.get("severity", "unknown")
|
|
633
|
-
call_type = decision.get("call_type")
|
|
634
|
-
|
|
635
|
-
try:
|
|
636
|
-
# Show a blocking dialog and wait for user decision
|
|
637
|
-
confirmation_request = ConfirmationRequest(
|
|
638
|
-
is_request=is_request,
|
|
639
|
-
tool_name=tool_name,
|
|
640
|
-
policy_reasons=policy_reasons,
|
|
641
|
-
content_data=content_data,
|
|
642
|
-
severity=severity,
|
|
643
|
-
event_id=event_id,
|
|
644
|
-
operation_type=operation_type,
|
|
645
|
-
server_name=self.wrapped_server_name,
|
|
646
|
-
timeout_seconds=60
|
|
647
|
-
)
|
|
648
|
-
|
|
649
|
-
response = UserConfirmationDialog(
|
|
650
|
-
self.logger, self.audit_logger
|
|
651
|
-
).request_blocking_confirmation(confirmation_request, prompt_id, call_type)
|
|
652
|
-
|
|
653
|
-
# If we got here, user chose "Allow Anyway"
|
|
654
|
-
self.logger.info(f"User chose to 'allow anyway' a blocked {confirmation_request.operation_type} "
|
|
655
|
-
f"operation for tool '{tool_name}' (event: {event_id})")
|
|
656
|
-
|
|
657
|
-
await self._record_user_confirmation(event_id, is_request, response.user_decision, prompt_id, call_type)
|
|
658
|
-
return
|
|
659
|
-
|
|
660
|
-
except UserConfirmationError as e:
|
|
661
|
-
# User chose to block or dialog failed
|
|
662
|
-
self.logger.warning(f"User blocking confirmation failed: {e}")
|
|
663
|
-
await self._record_user_confirmation(event_id, is_request, UserDecision.BLOCK, prompt_id, call_type)
|
|
664
|
-
reasons = "; ".join(policy_reasons)
|
|
665
|
-
raise error_class("Security Violation. User blocked the operation")
|
|
666
|
-
|
|
667
|
-
elif decision_type == "required_explicit_user_confirmation":
|
|
668
|
-
policy_reasons = decision.get("reasons", ["Security policy requires confirmation"])
|
|
669
|
-
severity = decision.get("severity", "unknown")
|
|
670
|
-
call_type = decision.get("call_type")
|
|
671
|
-
|
|
672
|
-
try:
|
|
673
|
-
confirmation_request = ConfirmationRequest(
|
|
674
|
-
is_request=is_request,
|
|
675
|
-
tool_name=tool_name,
|
|
676
|
-
policy_reasons=policy_reasons,
|
|
677
|
-
content_data=content_data,
|
|
678
|
-
severity=severity,
|
|
679
|
-
event_id=event_id,
|
|
680
|
-
operation_type=operation_type,
|
|
681
|
-
server_name=self.wrapped_server_name,
|
|
682
|
-
timeout_seconds=60
|
|
683
|
-
)
|
|
684
|
-
|
|
685
|
-
# only show YES_ALWAYS if call_type exists
|
|
686
|
-
options = DialogOptions(
|
|
687
|
-
show_always_allow=(call_type is not None),
|
|
688
|
-
show_always_block=False
|
|
689
|
-
)
|
|
690
|
-
|
|
691
|
-
response = UserConfirmationDialog(
|
|
692
|
-
self.logger, self.audit_logger
|
|
693
|
-
).request_confirmation(confirmation_request, prompt_id, call_type, options)
|
|
694
|
-
|
|
695
|
-
# If we got here, user approved the operation
|
|
696
|
-
self.logger.info(f"User {response.user_decision.value} {confirmation_request.operation_type} "
|
|
697
|
-
f"operation for tool '{tool_name}' (event: {event_id})")
|
|
698
|
-
|
|
699
|
-
await self._record_user_confirmation(event_id, is_request, response.user_decision, prompt_id, call_type)
|
|
700
|
-
return
|
|
701
|
-
|
|
702
|
-
except UserConfirmationError as e:
|
|
703
|
-
# User denied confirmation or dialog failed
|
|
704
|
-
self.logger.warning(f"User confirmation failed: {e}")
|
|
705
|
-
await self._record_user_confirmation(event_id, is_request, UserDecision.BLOCK, prompt_id, call_type)
|
|
706
|
-
raise error_class("Security Violation. User blocked the operation")
|
|
707
|
-
|
|
708
|
-
elif decision_type == "need_more_info":
|
|
709
|
-
stage_title = 'CLIENT REQUEST' if is_request else 'TOOL RESPONSE'
|
|
710
|
-
|
|
711
|
-
# Create an actionable error message for the AI agent
|
|
712
|
-
reasons = decision.get("reasons", [])
|
|
713
|
-
need_fields = decision.get("need_fields", [])
|
|
714
|
-
|
|
715
|
-
error_parts = [
|
|
716
|
-
f"SECURITY POLICY NEEDS MORE INFORMATION FOR REVIEWING {stage_title}:",
|
|
717
|
-
'\n'.join(reasons),
|
|
718
|
-
'' # newline
|
|
719
|
-
]
|
|
720
|
-
|
|
721
|
-
if need_fields:
|
|
722
|
-
# Convert server field names to wrapper field names for the AI agent
|
|
723
|
-
wrapper_field_mapping = {
|
|
724
|
-
"context.agent.intent": "__wrapper_modelIntent",
|
|
725
|
-
"context.agent.plan": "__wrapper_modelPlan",
|
|
726
|
-
"context.agent.expectedOutputs": "__wrapper_modelExpectedOutputs",
|
|
727
|
-
"context.agent.user_prompt": "__wrapper_userPrompt",
|
|
728
|
-
"context.agent.user_prompt_id": "__wrapper_userPromptId",
|
|
729
|
-
"context.agent.context_summary": "__wrapper_contextSummary",
|
|
730
|
-
"context.workspace.current_files": "__wrapper_currentFiles",
|
|
731
|
-
}
|
|
732
|
-
|
|
733
|
-
missing_wrapper_fields = []
|
|
734
|
-
for field in need_fields:
|
|
735
|
-
wrapper_field = wrapper_field_mapping.get(field, field)
|
|
736
|
-
missing_wrapper_fields.append(wrapper_field)
|
|
737
|
-
|
|
738
|
-
if missing_wrapper_fields:
|
|
739
|
-
error_parts.append("AFFECTED FIELDS:")
|
|
740
|
-
error_parts.extend(missing_wrapper_fields)
|
|
741
|
-
else:
|
|
742
|
-
error_parts.append("MISSING INFORMATION:")
|
|
743
|
-
error_parts.extend(need_fields)
|
|
744
|
-
|
|
745
|
-
|
|
746
|
-
error_parts.append("\nMANDATORY ACTIONS:")
|
|
747
|
-
error_parts.append("1. Add/Edit ALL affected fields according to the required information")
|
|
748
|
-
error_parts.append("2. Retry the tool call")
|
|
749
|
-
|
|
750
|
-
actionable_message = "\n".join(error_parts)
|
|
751
|
-
raise error_class(actionable_message)
|