feat: add secure agent and MCP extensions
This commit is contained in:
hailin
354
backend/apps/tools/tests/test_tool_runtime.py
Normal file
354
backend/apps/tools/tests/test_tool_runtime.py
Normal file
@@ -0,0 +1,354 @@
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from django.contrib.auth import get_user_model
|
||||
from django.db import transaction
|
||||
from django.test import override_settings
|
||||
|
||||
from apps.agents.models import AgentRun
|
||||
from apps.agents.tests.helpers import create_agent
|
||||
from apps.audit.models import OperationAuditLog
|
||||
from apps.modelhub.crypto import encrypt_secret
|
||||
from apps.tickets.models import Ticket
|
||||
from apps.tools.models import (
|
||||
ToolAuthorization,
|
||||
ToolCall,
|
||||
ToolCallStatus,
|
||||
ToolProvider,
|
||||
ToolRiskLevel,
|
||||
UnifiedTool,
|
||||
)
|
||||
from apps.tools.services import (
|
||||
ToolExecutionError,
|
||||
decide_tool_confirmation,
|
||||
request_tool_call,
|
||||
sync_mcp_tools,
|
||||
)
|
||||
from apps.workitems.models import WorkItem, WorkItemType
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, payload):
|
||||
self.payload = payload
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
return None
|
||||
|
||||
def read(self):
|
||||
return json.dumps(self.payload).encode()
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_http_mcp_sync_and_read_only_call_are_audited_and_idempotent(
|
||||
django_capture_on_commit_callbacks,
|
||||
monkeypatch,
|
||||
):
|
||||
responses = iter(
|
||||
[
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "ai-ticket-platform",
|
||||
"result": {
|
||||
"tools": [
|
||||
{
|
||||
"name": "lookup_customer",
|
||||
"description": "只读查询客户状态",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {"customer_id": {"type": "string"}},
|
||||
"required": ["customer_id"],
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
},
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "ai-ticket-platform",
|
||||
"result": {"content": [{"type": "text", "text": "active"}]},
|
||||
},
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"apps.tools.services.urlopen",
|
||||
lambda *args, **kwargs: FakeResponse(next(responses)),
|
||||
)
|
||||
provider = ToolProvider.objects.create(
|
||||
name="内部查询 MCP",
|
||||
base_url="https://mcp.internal.example/rpc",
|
||||
)
|
||||
tools = sync_mcp_tools(provider)
|
||||
assert len(tools) == 1
|
||||
tool = tools[0]
|
||||
assert tool.code == f"mcp.{provider.id}.lookup_customer"
|
||||
actor = get_user_model().objects.create_superuser(
|
||||
username="tool-admin",
|
||||
email="tool-admin@example.com",
|
||||
password="password",
|
||||
)
|
||||
|
||||
with django_capture_on_commit_callbacks(execute=True):
|
||||
first = request_tool_call(
|
||||
tool,
|
||||
{"customer_id": "C-001"},
|
||||
user=actor,
|
||||
idempotency_key="lookup-C-001",
|
||||
)
|
||||
first.refresh_from_db()
|
||||
second = request_tool_call(
|
||||
tool,
|
||||
{"customer_id": "C-001"},
|
||||
user=actor,
|
||||
idempotency_key="lookup-C-001",
|
||||
)
|
||||
|
||||
assert first.status == ToolCallStatus.SUCCEEDED
|
||||
assert first.output_payload["content"][0]["text"] == "active"
|
||||
assert second.id == first.id
|
||||
assert ToolCall.objects.count() == 1
|
||||
assert OperationAuditLog.objects.filter(
|
||||
module="tools",
|
||||
action="tool_call_finished",
|
||||
target_id=str(first.public_id),
|
||||
).exists()
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_high_risk_tool_requires_independent_confirmation(
|
||||
django_capture_on_commit_callbacks,
|
||||
monkeypatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
"apps.tools.services.urlopen",
|
||||
lambda *args, **kwargs: FakeResponse(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "ai-ticket-platform",
|
||||
"result": {"content": [{"type": "text", "text": "done"}]},
|
||||
}
|
||||
),
|
||||
)
|
||||
provider = ToolProvider.objects.create(
|
||||
name="变更 MCP",
|
||||
base_url="https://mcp.internal.example/rpc",
|
||||
)
|
||||
tool = UnifiedTool.objects.create(
|
||||
code="mcp.change.restart",
|
||||
provider=provider,
|
||||
external_name="restart_service",
|
||||
name="重启服务",
|
||||
input_schema={"type": "object"},
|
||||
risk_level=ToolRiskLevel.HIGH_RISK,
|
||||
)
|
||||
requester = get_user_model().objects.create_superuser(
|
||||
username="tool-requester",
|
||||
email="requester@example.com",
|
||||
password="password",
|
||||
)
|
||||
reviewer = get_user_model().objects.create_superuser(
|
||||
username="tool-reviewer",
|
||||
email="reviewer@example.com",
|
||||
password="password",
|
||||
)
|
||||
call = request_tool_call(tool, {}, user=requester)
|
||||
assert call.status == ToolCallStatus.PENDING_CONFIRMATION
|
||||
assert call.confirmation.status == "pending"
|
||||
task = WorkItem.objects.get(task_type=WorkItemType.TOOL_CONFIRMATION)
|
||||
assert task.status == "active"
|
||||
|
||||
with pytest.raises(ToolExecutionError, match="不能由申请人自审"):
|
||||
decide_tool_confirmation(call, reviewer=requester, decision="approve")
|
||||
|
||||
with django_capture_on_commit_callbacks(execute=True):
|
||||
decide_tool_confirmation(
|
||||
call,
|
||||
reviewer=reviewer,
|
||||
decision="approve",
|
||||
reason="已核对变更窗口",
|
||||
)
|
||||
call.refresh_from_db()
|
||||
call.confirmation.refresh_from_db()
|
||||
assert call.status == ToolCallStatus.SUCCEEDED
|
||||
assert call.confirmation.status == "approved"
|
||||
assert call.confirmation.decided_by == reviewer
|
||||
task.refresh_from_db()
|
||||
assert task.status == "completed"
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_tool_api_hides_credentials_and_blocks_unprivileged_invocation(
|
||||
admin_client,
|
||||
client,
|
||||
):
|
||||
response = admin_client.post(
|
||||
"/api/tools/providers",
|
||||
{
|
||||
"name": "受控 MCP",
|
||||
"provider_type": "mcp_http",
|
||||
"base_url": "https://mcp.internal.example/rpc",
|
||||
"token": "top-secret-token",
|
||||
},
|
||||
content_type="application/json",
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["credential_set"] is True
|
||||
assert "token" not in response.json()
|
||||
|
||||
user = get_user_model().objects.create_user(
|
||||
username="unprivileged-tool-user",
|
||||
password="password",
|
||||
)
|
||||
client.force_login(user)
|
||||
response = client.post(
|
||||
"/api/tools/calls",
|
||||
{"tool_code": "missing", "input": {}},
|
||||
content_type="application/json",
|
||||
)
|
||||
assert response.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
def test_tool_governance_api_lists_grants_and_calls(admin_client):
|
||||
provider = ToolProvider.objects.create(
|
||||
name="治理 MCP",
|
||||
base_url="https://mcp.internal.example/rpc",
|
||||
)
|
||||
tool = UnifiedTool.objects.create(
|
||||
code="mcp.governance.lookup",
|
||||
provider=provider,
|
||||
external_name="lookup",
|
||||
name="治理查询",
|
||||
input_schema={"type": "object"},
|
||||
)
|
||||
|
||||
response = admin_client.post(
|
||||
f"/api/tools/{tool.id}/authorizations",
|
||||
{"principal_type": "group", "principal_id": 7, "allowed": True},
|
||||
content_type="application/json",
|
||||
)
|
||||
assert response.status_code == 200
|
||||
response = admin_client.get(f"/api/tools/{tool.id}/authorizations")
|
||||
assert response.status_code == 200
|
||||
assert response.json()[0]["principal_type"] == "group"
|
||||
|
||||
response = admin_client.patch(
|
||||
f"/api/tools/providers/{provider.id}",
|
||||
{"timeout_seconds": 45, "enabled": False},
|
||||
content_type="application/json",
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["timeout_seconds"] == 45
|
||||
assert response.json()["enabled"] is False
|
||||
|
||||
actor = get_user_model().objects.filter(is_superuser=True).first()
|
||||
call = ToolCall.objects.create(
|
||||
tool=tool,
|
||||
requested_by=actor,
|
||||
status=ToolCallStatus.SUCCEEDED,
|
||||
input_payload={"query": "status"},
|
||||
output_payload={"status": "ok"},
|
||||
)
|
||||
response = admin_client.get("/api/tools/calls")
|
||||
assert response.status_code == 200
|
||||
payload = response.json()
|
||||
assert payload[0]["public_id"] == str(call.public_id)
|
||||
assert payload[0]["requested_by_id"] == actor.id
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
@override_settings(INTERNAL_MCP_TOKEN="internal-ticket-token")
|
||||
def test_internal_ticket_mcp_requires_bearer_and_queries_real_tickets(client):
|
||||
Ticket.objects.create(
|
||||
ticket_no="TK-MCP-001",
|
||||
title="客户无法登录控制台",
|
||||
description="登录后立即返回认证失败",
|
||||
status="pending",
|
||||
)
|
||||
rpc = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": "gateway-test",
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": "search_tickets",
|
||||
"arguments": {"query": "登录", "limit": 5},
|
||||
},
|
||||
}
|
||||
unauthorized = client.post(
|
||||
"/internal-mcp",
|
||||
rpc,
|
||||
content_type="application/json",
|
||||
)
|
||||
assert unauthorized.status_code == 401
|
||||
|
||||
response = client.post(
|
||||
"/internal-mcp",
|
||||
rpc,
|
||||
content_type="application/json",
|
||||
HTTP_AUTHORIZATION="Bearer internal-ticket-token",
|
||||
)
|
||||
assert response.status_code == 200
|
||||
structured = response.json()["result"]["structuredContent"]
|
||||
assert structured["count"] == 1
|
||||
assert structured["tickets"][0]["ticket_no"] == "TK-MCP-001"
|
||||
|
||||
|
||||
@pytest.mark.django_db(transaction=True)
|
||||
@override_settings(INTERNAL_MCP_TOKEN="real-http-ticket-token")
|
||||
def test_internal_ticket_mcp_real_http_agent_call_end_to_end(live_server):
|
||||
ticket = Ticket.objects.create(
|
||||
ticket_no="TK-MCP-E2E",
|
||||
title="生产订单同步延迟",
|
||||
description="订单同步任务积压，需要查询同类工单",
|
||||
status="in_progress",
|
||||
)
|
||||
actor = get_user_model().objects.create_superuser(
|
||||
username="mcp-e2e-user",
|
||||
email="mcp-e2e@example.com",
|
||||
password="password",
|
||||
)
|
||||
definition, version = create_agent(published=True, two_steps=False)
|
||||
provider = ToolProvider.objects.create(
|
||||
name="平台内部工单 MCP",
|
||||
base_url=f"{live_server.url}/internal-mcp",
|
||||
token_encrypted=encrypt_secret("real-http-ticket-token"),
|
||||
)
|
||||
tools = sync_mcp_tools(provider)
|
||||
assert len(tools) == 1
|
||||
tool = tools[0]
|
||||
assert tool.external_name == "search_tickets"
|
||||
version.allowed_actions = [f"tool.call:{tool.code}"]
|
||||
version.save(update_fields=["allowed_actions"])
|
||||
run = AgentRun.objects.create(
|
||||
definition=definition,
|
||||
agent_version=version,
|
||||
created_by=actor,
|
||||
config_snapshot={
|
||||
"agent": {"allowed_actions": [f"tool.call:{tool.code}"]}
|
||||
},
|
||||
)
|
||||
ToolAuthorization.objects.create(
|
||||
tool=tool,
|
||||
principal_type="agent",
|
||||
principal_id=definition.id,
|
||||
allowed=True,
|
||||
)
|
||||
|
||||
with transaction.atomic():
|
||||
call = request_tool_call(
|
||||
tool,
|
||||
{"query": "订单同步", "status": "in_progress"},
|
||||
user=actor,
|
||||
run=run,
|
||||
idempotency_key="internal-ticket-mcp-e2e",
|
||||
)
|
||||
call.refresh_from_db()
|
||||
|
||||
assert call.status == ToolCallStatus.SUCCEEDED
|
||||
structured = call.output_payload["structuredContent"]
|
||||
assert structured["count"] == 1
|
||||
assert structured["tickets"][0]["id"] == ticket.id
|
||||
assert call.authorization_snapshot["agent_grant"] is True
|
||||
assert call.risk_snapshot["risk_level"] == ToolRiskLevel.READ_ONLY
|
||||
Reference in New Issue
Block a user