Class: Aidp::Security::McpToolRiskClassifier

Inherits:
Object
  • Object
show all
Defined in:
lib/aidp/security/mcp_tool_risk_classifier.rb

Overview

Generates an MCP tool risk profile using an AI call at configuration time.

Constant Summary collapse

RISK_FLAGS = 
%w[untrusted_input private_data egress].freeze
RISK_LEVELS = 
%w[low medium high].freeze
GENERATION_PROMPT = 
<<~PROMPT
  You are classifying MCP tools for Rule-of-Two security enforcement.

  For each MCP tool, determine which risk flags should be enabled when that
  tool is available to an agent:
  - untrusted_input: processes or imports untrusted external/user-controlled content
  - private_data: can access secrets, credentials, local files, repos, databases, or other sensitive data
  - egress: can communicate externally, push data, call remote services, or send messages off-machine

  Also assign a risk_level of low, medium, or high based on the combined impact.

  Tools to classify:
  {{tools_json}}

  Respond with ONLY valid JSON in this format:
  {
    "tools": [
      {
        "name": "filesystem",
        "flags": ["private_data"],
        "risk_level": "medium",
        "rationale": "Can read and write local files that may contain secrets."
      }
    ]
  }
PROMPT

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(config, project_dir:, provider_factory: nil, provider_info_class: nil, time_source: Time) ⇒ McpToolRiskClassifier

Returns a new instance of McpToolRiskClassifier.



46
47
48
49
50
51
52
 
# File 'lib/aidp/security/mcp_tool_risk_classifier.rb', line 46

def initialize(config, project_dir:, provider_factory: nil, provider_info_class: nil, time_source: Time)
  @config = config
  @project_dir = project_dir
  @provider_factory = provider_factory || Aidp::Harness::ProviderFactory.new(config)
  @provider_info_class = provider_info_class || Aidp::Harness::ProviderInfo
  @time_source = time_source
end

Instance Attribute Details

#config ⇒ Object (readonly)

Returns the value of attribute config.



44
45
46
 
# File 'lib/aidp/security/mcp_tool_risk_classifier.rb', line 44

def config
  @config
end

#project_dir ⇒ Object (readonly)

Returns the value of attribute project_dir.



44
45
46
 
# File 'lib/aidp/security/mcp_tool_risk_classifier.rb', line 44

def project_dir
  @project_dir
end

#provider_factory ⇒ Object (readonly)

Returns the value of attribute provider_factory.



44
45
46
 
# File 'lib/aidp/security/mcp_tool_risk_classifier.rb', line 44

def provider_factory
  @provider_factory
end

Instance Method Details

#collect_tools(providers: nil, force_refresh: false, force_refresh_providers: nil) ⇒ Object 



74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
 
# File 'lib/aidp/security/mcp_tool_risk_classifier.rb', line 74

def collect_tools(providers: nil, force_refresh: false, force_refresh_providers: nil)
  provider_names = Array(providers || config.configured_providers).map(&:to_s)
  force_refresh_provider_names = Array(force_refresh_providers).map(&:to_s)

  provider_names.each_with_object({}) do |provider_name, tools|
    refresh_provider = refresh_provider?(provider_name, force_refresh, force_refresh_provider_names)
    provider_info = @provider_info_class.new(provider_name, project_dir)
    info = provider_info.info(force_refresh: refresh_provider)
    validate_refreshed_provider_info!(provider_name, info, refresh_provider)
    next unless info&.dig(:mcp_support)

    enabled_mcp_servers(info).each do |server|
      next if server_name(server).empty?

      entry = (tools[server_name(server)] ||= {
        name: server_name(server),
        descriptions: [],
        providers: []
      })
      description = server[:description].to_s.strip
      entry[:descriptions] << description unless description.empty?
      entry[:providers] << provider_name
    end
  end.values.map do |tool|
    tool[:descriptions].uniq!
    tool[:providers].uniq!
    tool
  end.sort_by { |tool| tool[:name] }
end

#generate!(providers: nil, tier: "mini", force_refresh: false, force_refresh_providers: nil) ⇒ Object 



54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
 
# File 'lib/aidp/security/mcp_tool_risk_classifier.rb', line 54

def generate!(providers: nil, tier: "mini", force_refresh: false, force_refresh_providers: nil)
  tools = collect_tools(
    providers: providers,
    force_refresh: force_refresh,
    force_refresh_providers: force_refresh_providers
  )
  return write_empty_profile if tools.empty?

  provider_name, model_name = select_model(tier)
  response = call_ai(provider_name, model_name, build_prompt(tools))
  profile = build_profile(response, model_name, expected_tools: tools)
  profile.save!(project_dir)
  profile
rescue => e
  Aidp.log_error("security.mcp_classifier", "generation_failed",
    error: e.message,
    error_class: e.class.name)
  raise
end