Skip to content

Commit 826c82f

Browse files
gorkem-bwlgorkem-bwl
authored
Merge pull request #3688 from verifywise-ai/feature/mcp-gateway-operational-improvements
Feature: Mcp Gateway Improvements
2 parents ce0b892 + 435a111 commit 826c82f

File tree

11 files changed

+240
-32
lines changed

11 files changed

+240
-32
lines changed

‎AIGateway/src/crud/mcp_approvals.py‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,52 @@ async def get_approval_history(
129129
return []
130130

131131

132+
async def get_pending_request(org_id: int, agent_key_id: int, tool_name: str) -> Optional[dict]:
133+
"""Return an existing pending, non-expired approval request for this agent+tool."""
134+
async with get_db() as db:
135+
result = await db.execute(
136+
text("""
137+
SELECT id, status, expires_at
138+
FROM ai_gateway_mcp_approval_requests
139+
WHERE organization_id = :org_id
140+
AND agent_key_id = :agent_key_id
141+
AND tool_name = :tool_name
142+
AND status = 'pending'
143+
AND expires_at > NOW()
144+
ORDER BY created_at DESC
145+
LIMIT 1
146+
"""),
147+
{"org_id": org_id, "agent_key_id": agent_key_id, "tool_name": tool_name},
148+
)
149+
row = result.mappings().first()
150+
if row is None:
151+
return None
152+
return dict(row)
153+
154+
155+
async def get_approved_request(org_id: int, agent_key_id: int, tool_name: str) -> Optional[dict]:
156+
"""Check if an approved, non-expired approval request exists for this agent+tool."""
157+
async with get_db() as db:
158+
result = await db.execute(
159+
text("""
160+
SELECT id, status, expires_at
161+
FROM ai_gateway_mcp_approval_requests
162+
WHERE organization_id = :org_id
163+
AND agent_key_id = :agent_key_id
164+
AND tool_name = :tool_name
165+
AND status = 'approved'
166+
AND expires_at > NOW()
167+
ORDER BY decided_at DESC
168+
LIMIT 1
169+
"""),
170+
{"org_id": org_id, "agent_key_id": agent_key_id, "tool_name": tool_name},
171+
)
172+
row = result.mappings().first()
173+
if row is None:
174+
return None
175+
return dict(row)
176+
177+
132178
async def get_approval_status(org_id: int, request_id: int) -> Optional[dict]:
133179
async with get_db() as db:
134180
result = await db.execute(
@@ -202,3 +248,15 @@ async def decide_approval(
202248
return None
203249
return dict(row)
204250
return None
251+
252+
253+
async def delete_expired_approval_requests(retention_days: int = 30, batch_size: int = 5000) -> int:
254+
"""Delete decided or expired approval requests older than retention_days in batches."""
255+
from utils.batch_delete import batch_delete_expired
256+
257+
return await batch_delete_expired(
258+
table="ai_gateway_mcp_approval_requests",
259+
where_clause="(status != 'pending' OR expires_at < NOW())",
260+
retention_days=retention_days,
261+
batch_size=batch_size,
262+
)

‎AIGateway/src/crud/mcp_audit.py‎

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -183,17 +183,13 @@ async def get_audit_stats_by_agent(org_id: int, days: int = 7) -> list[dict]:
183183
]
184184

185185

186-
async def delete_expired_audit_logs(retention_days: int = 30) -> int:
187-
"""Delete audit logs older than retention_days. Returns count of deleted rows."""
188-
retention_days = int(retention_days)
189-
190-
async with get_db() as db:
191-
result = await db.execute(
192-
text(f"""
193-
DELETE FROM ai_gateway_mcp_audit_logs
194-
WHERE created_at < NOW() - INTERVAL '{retention_days} days'
195-
RETURNING id
196-
"""),
197-
)
198-
await db.commit()
199-
return len(result.fetchall())
186+
async def delete_expired_audit_logs(retention_days: int = 30, batch_size: int = 5000) -> int:
187+
"""Delete audit logs older than retention_days in batches. Returns count of deleted rows."""
188+
from utils.batch_delete import batch_delete_expired
189+
190+
return await batch_delete_expired(
191+
table="ai_gateway_mcp_audit_logs",
192+
where_clause="",
193+
retention_days=retention_days,
194+
batch_size=batch_size,
195+
)

‎AIGateway/src/routers/mcp_audit.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,3 +106,18 @@ async def cleanup_audit_logs(request: Request):
106106

107107
deleted = await delete_expired_audit_logs(settings.mcp_audit_retention_days)
108108
return {"status": "success", "deleted": deleted}
109+
110+
111+
# ---------------------------------------------------------------------------
112+
# POST /mcp/audit/cleanup-approvals
113+
# ---------------------------------------------------------------------------
114+
115+
@router.post("/cleanup-approvals", status_code=status.HTTP_200_OK)
116+
async def cleanup_approval_requests(request: Request):
117+
"""Delete decided/expired approval requests older than retention period."""
118+
verify_internal_key(request)
119+
from config import settings
120+
from crud.mcp_approvals import delete_expired_approval_requests
121+
122+
deleted = await delete_expired_approval_requests(settings.mcp_audit_retention_days)
123+
return {"status": "success", "deleted": deleted}

‎AIGateway/src/routers/mcp_proxy.py‎

Lines changed: 49 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,13 @@
1212
import asyncio
1313
import logging
1414
import time
15+
from datetime import datetime, timezone, timedelta
1516

1617
from fastapi import APIRouter, Request, HTTPException
1718
from fastapi.responses import JSONResponse, StreamingResponse
1819

20+
from config import settings
21+
from crud.mcp_approvals import create_approval_request, get_approval_status, get_approved_request, get_pending_request
1922
from crud.mcp_tools import get_all_tools
2023
from services.mcp_audit_service import log_tool_call
2124
from services.mcp_guardrail_service import scan_tool_input
@@ -34,8 +37,11 @@
3437
router = APIRouter()
3538

3639

37-
def _jsonrpc_error(id, code: int, message: str) -> dict:
38-
return {"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message}}
40+
def _jsonrpc_error(id, code: int, message: str, data: dict | None = None) -> dict:
41+
error = {"code": code, "message": message}
42+
if data:
43+
error["data"] = data
44+
return {"jsonrpc": "2.0", "id": id, "error": error}
3945

4046

4147
def _jsonrpc_result(id, result: dict) -> dict:
@@ -135,7 +141,28 @@ async def _audit(status: str, summary: str | None, is_error: bool):
135141
await enforce_mcp_rate_limits(agent_key, tool_name)
136142

137143
if tool.get("requires_approval"):
138-
return JSONResponse(content=_jsonrpc_error(msg_id, -32001, "Tool requires approval"), status_code=200)
144+
# Check if an approved request already exists for this agent+tool
145+
approved = await get_approved_request(org_id, agent_key["id"], tool_name)
146+
if not approved:
147+
# Reuse an existing pending request if one exists, otherwise create a new one
148+
pending = await get_pending_request(org_id, agent_key["id"], tool_name)
149+
if pending:
150+
approval = pending
151+
else:
152+
expires_at = datetime.now(timezone.utc) + timedelta(seconds=settings.mcp_approval_expiry_seconds)
153+
approval = await create_approval_request(org_id, {
154+
"agent_key_id": agent_key["id"],
155+
"tool_id": tool.get("id"),
156+
"tool_name": tool_name,
157+
"arguments": arguments,
158+
"expires_at": expires_at,
159+
})
160+
await _audit("approval_required", f"Approval request {approval.get('id')} created", False)
161+
return JSONResponse(content=_jsonrpc_error(msg_id, -32001, "Tool requires approval", {
162+
"approval_id": approval.get("id"),
163+
"poll_endpoint": f"/v1/mcp/approvals/{approval.get('id')}/status",
164+
"expires_at": approval["expires_at"].isoformat() if hasattr(approval["expires_at"], "isoformat") else str(approval["expires_at"]),
165+
}), status_code=200)
139166

140167
scan_result = await scan_tool_input(org_id, tool_name, arguments)
141168
if scan_result and scan_result.blocked:
@@ -182,6 +209,25 @@ async def _audit(status: str, summary: str | None, is_error: bool):
182209
return JSONResponse(content=_jsonrpc_error(msg_id, -32601, f"Method not found: {method}"), status_code=200)
183210

184211

212+
@router.get("/v1/mcp/approvals/{request_id}/status")
213+
async def mcp_approval_status(request: Request, request_id: int):
214+
"""Poll approval status — authenticated via agent key."""
215+
agent_key = await _extract_agent_key(request)
216+
org_id = agent_key["organization_id"]
217+
218+
approval = await get_approval_status(org_id, request_id)
219+
if not approval:
220+
raise HTTPException(status_code=404, detail="Approval request not found")
221+
222+
return {
223+
"approval_id": approval["id"],
224+
"status": approval["status"],
225+
"decided_at": approval["decided_at"].isoformat() if approval.get("decided_at") else None,
226+
"decision_reason": approval.get("decision_reason"),
227+
"expires_at": approval["expires_at"].isoformat() if approval.get("expires_at") else None,
228+
}
229+
230+
185231
@router.get("/v1/mcp")
186232
async def mcp_sse(request: Request):
187233
"""SSE endpoint for server-initiated messages. Keep-alive for v1."""

‎AIGateway/src/services/mcp_audit_service.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,19 @@
1111

1212
logger = logging.getLogger("uvicorn")
1313

14+
MAX_ARGUMENTS_SIZE = 10_240 # 10 KB
15+
16+
17+
def _cap_arguments(arguments: dict | None) -> str:
18+
"""Serialize arguments to JSON, truncating to MAX_ARGUMENTS_SIZE."""
19+
if not arguments:
20+
return "{}"
21+
serialized = json.dumps(arguments)
22+
if len(serialized) <= MAX_ARGUMENTS_SIZE:
23+
return serialized
24+
original_size = len(serialized)
25+
return json.dumps({"_truncated": True, "_original_size": original_size})
26+
1427

1528
async def log_tool_call(
1629
organization_id: int,
@@ -48,7 +61,7 @@ async def log_tool_call(
4861
"agent_key_id": agent_key_id,
4962
"server_id": server_id,
5063
"tool_name": tool_name,
51-
"arguments": json.dumps(arguments) if arguments else "{}",
64+
"arguments": _cap_arguments(arguments),
5265
"result_status": result_status,
5366
"result_summary": result_summary,
5467
"is_error": is_error,

‎AIGateway/src/utils/batch_delete.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
"""Shared batched-delete helper for retention cleanup jobs."""
2+
3+
from sqlalchemy import text
4+
5+
from database.db import get_db
6+
7+
8+
async def batch_delete_expired(
9+
table: str,
10+
where_clause: str,
11+
retention_days: int = 30,
12+
batch_size: int = 5000,
13+
) -> int:
14+
"""Delete rows matching where_clause older than retention_days in batches.
15+
16+
``where_clause`` is appended after
17+
``WHERE created_at < NOW() - INTERVAL '<retention_days> days'``.
18+
Pass an empty string if no extra conditions are needed.
19+
20+
Returns total number of deleted rows.
21+
"""
22+
retention_days = int(retention_days)
23+
extra = f" AND {where_clause}" if where_clause else ""
24+
total_deleted = 0
25+
26+
async with get_db() as db:
27+
while True:
28+
result = await db.execute(
29+
text(f"""
30+
DELETE FROM {table}
31+
WHERE id IN (
32+
SELECT id FROM {table}
33+
WHERE created_at < NOW() - INTERVAL '{retention_days} days'
34+
{extra}
35+
LIMIT :batch_size
36+
)
37+
RETURNING id
38+
"""),
39+
{"batch_size": batch_size},
40+
)
41+
deleted = len(result.fetchall())
42+
await db.commit()
43+
total_deleted += deleted
44+
if deleted < batch_size:
45+
break
46+
47+
return total_deleted

‎Clients/src/presentation/components/breadcrumbs/routeMapping.ts‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,15 @@ export const routeMapping: Record<string, string> = {
150150
"/ai-gateway/settings/guardrails": "Guardrail settings",
151151
"/ai-gateway/settings/risks": "Suggested risks",
152152

153+
// MCP Gateway
154+
"/ai-gateway/mcp": "MCP Gateway",
155+
"/ai-gateway/mcp/agent-keys": "Agent keys",
156+
"/ai-gateway/mcp/servers": "Servers",
157+
"/ai-gateway/mcp/tools": "Tool catalog",
158+
"/ai-gateway/mcp/audit": "Audit log",
159+
"/ai-gateway/mcp/approvals": "Approvals",
160+
"/ai-gateway/mcp/guardrails": "Guardrails",
161+
153162
// Shadow AI
154163
"/shadow-ai": "Shadow AI",
155164
"/shadow-ai/insights": "Insights",
@@ -301,6 +310,15 @@ export const routeIconMapping: Record<string, () => React.ReactNode> = {
301310
"/ai-gateway/settings/guardrails": () => React.createElement(Settings, { size: 14, strokeWidth: 1.5 }),
302311
"/ai-gateway/settings/risks": () => React.createElement(Settings, { size: 14, strokeWidth: 1.5 }),
303312

313+
// MCP Gateway
314+
"/ai-gateway/mcp": () => React.createElement(Router, { size: 14, strokeWidth: 1.5 }),
315+
"/ai-gateway/mcp/agent-keys": () => React.createElement(KeyRound, { size: 14, strokeWidth: 1.5 }),
316+
"/ai-gateway/mcp/servers": () => React.createElement(Router, { size: 14, strokeWidth: 1.5 }),
317+
"/ai-gateway/mcp/tools": () => React.createElement(Layers, { size: 14, strokeWidth: 1.5 }),
318+
"/ai-gateway/mcp/audit": () => React.createElement(FileSearch, { size: 14, strokeWidth: 1.5 }),
319+
"/ai-gateway/mcp/approvals": () => React.createElement(ShieldCheck, { size: 14, strokeWidth: 1.5 }),
320+
"/ai-gateway/mcp/guardrails": () => React.createElement(ShieldAlert, { size: 14, strokeWidth: 1.5 }),
321+
304322
// Super Admin
305323
"/super-admin": () => React.createElement(Shield, { size: 14, strokeWidth: 1.5 }),
306324
"/super-admin/users": () => React.createElement(Users, { size: 14, strokeWidth: 1.5 }),

‎Clients/src/presentation/types/interfaces/i.header.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ export interface PageHeaderExtendedProps {
88
tipBoxEntity?: string;
99
summaryCards?: ReactNode;
1010
summaryCardsJoyrideId?: string;
11-
children: ReactNode;
11+
children?: ReactNode;
1212
alert?: ReactNode;
1313
loadingToast?: ReactNode;
1414
titleFontFamily?: string;

‎Servers/jobs/producer.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ export * from "../services/slack/slackProducer";
22

33
import { scheduleDailyNotification } from "../services/slack/slackProducer";
44
import logger from "../utils/logger/fileLogger";
5-
import { scheduleReportNotification, scheduleVendorReviewDateNotification, schedulePMMHourlyCheck, scheduleShadowAiJobs, schedulePolicyDueSoonNotification, scheduleAgentDiscoverySync, scheduleAiDetectionScanCheck, scheduleAiGatewayRiskDetection, scheduleAiGatewayCacheCleanup } from "../services/automations/automationProducer";
5+
import { scheduleReportNotification, scheduleVendorReviewDateNotification, schedulePMMHourlyCheck, scheduleShadowAiJobs, schedulePolicyDueSoonNotification, scheduleAgentDiscoverySync, scheduleAiDetectionScanCheck, scheduleAiGatewayRiskDetection, scheduleAiGatewayCacheCleanup, scheduleMcpGatewayCleanup } from "../services/automations/automationProducer";
66

77
export async function addAllJobs(): Promise<void> {
88
await scheduleDailyNotification();
@@ -15,6 +15,7 @@ export async function addAllJobs(): Promise<void> {
1515
await scheduleAiDetectionScanCheck();
1616
await scheduleAiGatewayRiskDetection();
1717
await scheduleAiGatewayCacheCleanup();
18+
await scheduleMcpGatewayCleanup();
1819
}
1920

2021
if (require.main === module) {

‎Servers/services/automations/automationProducer.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,3 +192,17 @@ export async function scheduleAiGatewayCacheCleanup() {
192192
},
193193
);
194194
}
195+
196+
export async function scheduleMcpGatewayCleanup() {
197+
logger.info("Adding MCP Gateway cleanup jobs to the queue...");
198+
// Daily at 3 AM — purge expired audit logs and decided approval requests
199+
await automationQueue.add(
200+
"mcp_audit_cleanup",
201+
{ type: "mcp_gateway" },
202+
{
203+
repeat: { pattern: "0 3 * * *" },
204+
removeOnComplete: true,
205+
removeOnFail: false,
206+
},
207+
);
208+
}

0 commit comments

Comments
 (0)