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.

Files changed (39) hide show
  1. ide_tools/__init__.py +12 -0
  2. ide_tools/common/__init__.py +5 -0
  3. ide_tools/common/hooks/__init__.py +5 -0
  4. ide_tools/common/hooks/init.py +124 -0
  5. ide_tools/common/hooks/output.py +63 -0
  6. ide_tools/common/hooks/prompt_submit.py +133 -0
  7. ide_tools/common/hooks/read_file.py +167 -0
  8. ide_tools/common/hooks/shell_execution.py +255 -0
  9. ide_tools/common/hooks/shell_parser_bashlex.py +277 -0
  10. ide_tools/common/hooks/types.py +34 -0
  11. ide_tools/common/hooks/utils.py +286 -0
  12. ide_tools/cursor/__init__.py +11 -0
  13. ide_tools/cursor/constants.py +58 -0
  14. ide_tools/cursor/format.py +35 -0
  15. ide_tools/cursor/router.py +100 -0
  16. ide_tools/router.py +48 -0
  17. main.py +11 -4
  18. {mcpower_proxy-0.0.65.dist-info → mcpower_proxy-0.0.74.dist-info}/METADATA +4 -3
  19. mcpower_proxy-0.0.74.dist-info/RECORD +60 -0
  20. {mcpower_proxy-0.0.65.dist-info → mcpower_proxy-0.0.74.dist-info}/top_level.txt +1 -0
  21. modules/apis/security_policy.py +11 -6
  22. modules/decision_handler.py +219 -0
  23. modules/logs/audit_trail.py +16 -15
  24. modules/logs/logger.py +14 -18
  25. modules/redaction/gitleaks_rules.py +1 -1
  26. modules/redaction/pii_rules.py +0 -48
  27. modules/redaction/redactor.py +112 -107
  28. modules/ui/__init__.py +1 -1
  29. modules/ui/confirmation.py +0 -1
  30. modules/utils/cli.py +36 -6
  31. modules/utils/ids.py +55 -10
  32. modules/utils/json.py +3 -3
  33. wrapper/__version__.py +1 -1
  34. wrapper/middleware.py +135 -217
  35. wrapper/server.py +19 -11
  36. mcpower_proxy-0.0.65.dist-info/RECORD +0 -43
  37. {mcpower_proxy-0.0.65.dist-info → mcpower_proxy-0.0.74.dist-info}/WHEEL +0 -0
  38. {mcpower_proxy-0.0.65.dist-info → mcpower_proxy-0.0.74.dist-info}/entry_points.txt +0 -0
  39. {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
- fd = os.open(str(file_path), os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
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(f"{id_type.title()} file created by another process, reading (attempt {attempt + 1}/{max_attempts})")
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 = Path.home() / ".mcpower" / "uid"
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
@@ -3,4 +3,4 @@
3
3
  Wrapper MCP Server Version
4
4
  """
5
5
 
6
- __version__ = "0.0.65"
6
+ __version__ = "0.0.74"
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
- # Check workspace roots and re-initialize app_uid if workspace changed
91
- workspace_roots = await self._extract_workspace_roots(context)
92
- current_workspace_root = workspace_roots[0] if workspace_roots else str(Path.home() / ".mcpower")
93
- if current_workspace_root != self._last_workspace_root:
94
- self.logger.debug(f"Workspace root changed from {self._last_workspace_root} to {current_workspace_root}")
95
- self._last_workspace_root = current_workspace_root
96
- self.app_id = read_app_uid(logger=self.logger, project_folder_path=current_workspace_root)
97
- self.audit_logger.set_app_uid(self.app_id)
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
- call_next_callback = call_next
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, call_next)
118
- case "resources/list" | "resources/templates/list" | "prompts/list":
119
- return await call_next_callback(context)
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=call_next_callback,
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.info(f"PROFILE: {operation_type} id: {event_id} inspect_request duration: {on_inspect_request_duration:.2f} seconds")
246
+ self.logger.debug(
247
+ f"PROFILE: {operation_type} id: {event_id} inspect_request duration: {on_inspect_request_duration:.2f} seconds")
225
248
 
226
- await self._enforce_decision(
227
- decision=request_decision,
228
- error_class=error_class,
229
- base_message=f"{operation_type.title()} request blocked by security policy",
230
- is_request=True,
231
- event_id=event_id,
232
- tool_name=tool_name,
233
- content_data=tool_args,
234
- operation_type=operation_type,
235
- prompt_id=prompt_id
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.info(f"PROFILE: {operation_type} id: {event_id} call_next duration: {on_call_next_duration:.2f} seconds")
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.info(f"PROFILE: {operation_type} id: {event_id} inspect_response duration: {on_inspect_response_duration:.2f} seconds")
310
+ self.logger.debug(
311
+ f"PROFILE: {operation_type} id: {event_id} inspect_response duration: {on_inspect_response_duration:.2f} seconds")
279
312
 
280
- await self._enforce_decision(
281
- decision=response_decision,
282
- error_class=error_class,
283
- base_message=f"{operation_type.title()} response blocked by security policy",
284
- is_request=False,
285
- event_id=event_id,
286
- tool_name=tool_name,
287
- content_data=response_content,
288
- operation_type=operation_type,
289
- prompt_id=prompt_id
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.info(f"PROFILE: {operation_type} id: {event_id} duration: {on_handle_operation_duration:.2f} seconds")
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
- result = await call_next(context)
311
- self.logger.info(f"PROFILE: tools/list call_next duration: {time.time() - on_handle_tools_list_start_time:.2f} seconds id: {event_id}")
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.info(f"PROFILE: tools/list enhanced_result duration: {on_handle_tools_list_duration:.2f} seconds id: {event_id}")
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.info(f"PROFILE: tools/list result duration: {on_handle_tools_list_duration:.2f} seconds id: {event_id}")
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
- server_name=base_dict["server"]["name"],
508
- server_transport=base_dict["server"]["transport"],
509
- tool_name=base_dict["tool"]["name"] or base_dict["tool"]["method"],
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
- server_name=base_dict["server"]["name"],
536
- server_transport=base_dict["server"]["transport"],
537
- tool_name=base_dict["tool"]["name"] or base_dict["tool"]["method"],
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)