wlw/ai-agent
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: update MCP tool handling to treat administrator-selected tools as direct tools without confirmation

Browse Source
This commit is contained in:
mlogclub
2026-07-26 12:39:34 +08:00
parent 9d9c8a8346
commit c021c5568c
3 changed files with 17 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/ai/tooling/registry.go
+5 -4
View File
@@ -74,15 +74,16 @@ func (r *Registry) Resolve(toolCode string) (Definition, error) {
if serverCode == "" || toolName == "" { if serverCode == "" || toolName == "" {
return Definition{}, fmt.Errorf("unsupported tool code: %s", toolCode) return Definition{}, fmt.Errorf("unsupported tool code: %s", toolCode)
} }
// MCP metadata cannot reliably describe side effects. Treat it as sensitive // MCP tools are explicitly selected by an administrator before an Agent can
// until an administrator provides a more specific policy in a later phase. // call them. Treat that persisted allow-list as the authorization boundary;
// only tools with an explicit built-in policy require extra confirmation.
return Definition{ return Definition{
Code: toolCode, Code: toolCode,
Name: toolName, Name: toolName,
InputSchema: map[string]any{"type": "object", "additionalProperties": true}, InputSchema: map[string]any{"type": "object", "additionalProperties": true},
SourceType: enums.ToolSourceTypeMCP, SourceType: enums.ToolSourceTypeMCP,
RiskLevel: RiskLevelSensitive, RiskLevel: RiskLevelRead,
RequireConfirmation: true, RequireConfirmation: false,
MaxCallsPerRun: 3, MaxCallsPerRun: 3,
TimeoutMS: 30000, TimeoutMS: 30000,
IdempotencyMode: "caller", IdempotencyMode: "caller",
internal/ai/tooling/registry_test.go
+8 -7
View File
@@ -61,15 +61,15 @@ func TestRegistryIncludesAllTicketDraftToolInputs(t *testing.T) {
} }
} }
func TestRegistryTreatsMCPToolsAsSensitive(t *testing.T) { func TestRegistryTreatsAdministratorSelectedMCPToolsAsDirectTools(t *testing.T) {
definition, err := DefaultRegistry.Resolve("knowledge/search") definition, err := DefaultRegistry.Resolve("knowledge/search")
if err != nil { if err != nil {
t.Fatalf("Resolve returned error: %v", err) t.Fatalf("Resolve returned error: %v", err)
} }
if definition.RiskLevel != RiskLevelSensitive || !definition.RequireConfirmation { if definition.RiskLevel != RiskLevelRead || definition.RequireConfirmation {
t.Fatalf("unexpected MCP definition: %#v", definition) t.Fatalf("unexpected MCP definition: %#v", definition)
} }
if err := DefaultRegistry.Authorize(definition, Policy{AllowedToolCodes: []string{"knowledge/search"}, Confirmed: true}); err != nil { if err := DefaultRegistry.Authorize(definition, Policy{AllowedToolCodes: []string{"knowledge/search"}, AllowedRiskLevels: []string{RiskLevelRead}}); err != nil {
t.Fatalf("Authorize returned error: %v", err) t.Fatalf("Authorize returned error: %v", err)
} }
} }
@@ -91,13 +91,14 @@ func TestNormalizeCustomerReplyRejectsSecretAndNormalizesText(t *testing.T) {
} }
} }
func TestMCPExecutorRejectsUnconfirmedToolBeforeRuntimeCall(t *testing.T) { func TestMCPExecutorAllowsSelectedToolThroughAuthorization(t *testing.T) {
executor := NewMCPExecutor(DefaultRegistry, nil) executor := NewMCPExecutor(DefaultRegistry, nil)
_, _, err := executor.Execute(t.Context(), "knowledge/search", nil, Policy{ _, _, err := executor.Execute(t.Context(), "knowledge/search", nil, Policy{
AllowedToolCodes: []string{"knowledge/search"}, AllowedToolCodes: []string{"knowledge/search"},
AllowedRiskLevels: []string{RiskLevelRead},
}) })
if err == nil || !strings.Contains(err.Error(), "confirmation") { if err == nil || !strings.Contains(err.Error(), "runtime is not configured") {
t.Fatalf("expected confirmation rejection, got %v", err) t.Fatalf("expected authorization to pass before the missing runtime error, got %v", err)
} }
} }
internal/services/ai_agent_workflow_service_test.go
+4 -4
View File
@@ -336,11 +336,11 @@ func TestAIAgentServiceRejectsPublishWithUnavailableModelConfig(t *testing.T) {
} }
} }
func TestAIAgentServiceRejectsPublishWithSensitiveDirectTool(t *testing.T) { func TestAIAgentServiceAllowsPublishWithAdministratorSelectedMCPTool(t *testing.T) {
setupAIAgentWorkflowTestDB(t) setupAIAgentWorkflowTestDB(t)
operator := aiAgentWorkflowTestOperator() operator := aiAgentWorkflowTestOperator()
agent, err := AIAgentService.CreateAIAgent(request.CreateAIAgentRequest{ agent, err := AIAgentService.CreateAIAgent(request.CreateAIAgentRequest{
Name: "sensitive tool agent", AIConfigID: createAIAgentWorkflowTestConfig(t), RuntimeMode: enums.AIAgentRuntimeModeAutonomous, Name: "mcp tool agent", AIConfigID: createAIAgentWorkflowTestConfig(t), RuntimeMode: enums.AIAgentRuntimeModeAutonomous,
ServiceMode: enums.IMConversationServiceModeAIOnly, HandoffMode: enums.AIAgentHandoffModeWaitPool, FallbackMode: enums.AIAgentFallbackModeNoAnswer, ServiceMode: enums.IMConversationServiceModeAIOnly, HandoffMode: enums.AIAgentHandoffModeWaitPool, FallbackMode: enums.AIAgentFallbackModeNoAnswer,
}, operator) }, operator)
if err != nil { if err != nil {
@@ -349,8 +349,8 @@ func TestAIAgentServiceRejectsPublishWithSensitiveDirectTool(t *testing.T) {
if err := sqls.DB().Model(&models.AIAgent{}).Where("id = ?", agent.ID).Update("allowed_mcp_tools", `[{"toolCode":"mcp/demo/write_order"}]`).Error; err != nil { if err := sqls.DB().Model(&models.AIAgent{}).Where("id = ?", agent.ID).Update("allowed_mcp_tools", `[{"toolCode":"mcp/demo/write_order"}]`).Error; err != nil {
t.Fatalf("set direct tool: %v", err) t.Fatalf("set direct tool: %v", err)
} }
if _, err := AIAgentService.PublishAIAgent(agent.ID, operator); err == nil || !strings.Contains(err.Error(), "confirmed playbook") { if _, err := AIAgentService.PublishAIAgent(agent.ID, operator); err != nil {
t.Fatalf("expected sensitive direct tool publish rejection, got %v", err) t.Fatalf("expected administrator-selected MCP tool to be publishable, got %v", err)
} }
} }