Module: PWN::AI::Agent::Registry

Defined in:
lib/pwn/ai/agent/registry.rb

Overview

Central registry for pwn-ai agent tools.

Each file under lib/pwn/ai/agent/tools/*.rb calls PWN::AI::Agent::Registry.register(...) at load time to declare a JSON-Schema (what the LLM sees) and a handler lambda (what pwn runs).

Registry.definitions(...) returns the OpenAI-format tools: array; Registry.lookup(name:) returns the entry for dispatch.

Import chain (circular-import safe):

agent/registry.rb        (no deps on tool files)
     ^
agent/tools/*.rb         (require registry, call .register at top level)
     ^
agent/loop.rb            (calls Registry.discover then .definitions)

DYNAMIC TOOL-SET SLIMMING (local-model scaffolding)

Shipping all ~47 tool schemas on every call overwhelms a 35B local model — it mis-routes (extro_rf_tune for a git question) because the choice space is huge. When PWN::Env[:agent][:tool_router] is truthy (or nil while active==:ollama) AND definitions(relevance:) is passed, the pool is reduced to CORE_TOOLS + the top-K keyword-ranked matches. Routing accuracy is fed back into Metrics under name:'tool_router' so the router itself becomes a learned component.

Defined Under Namespace

Classes: Entry

Constant Summary collapse

CORE_TOOLS = 
%w[memory_recall session_recall skills_recall pwn_eval shell
mistakes_record mistakes_resolve learning_note_outcome
memory_remember skills_update].freeze
DEFAULT_PREFERENCE =

Schema order = CORE_TOOLS. Current session is injected (RECENT TURNS). Then memory / prior sessions / skills, then pwn_eval before shell. opts, or PWN::Env[:agent][:tool_preference]. Explicit nil/empty order disables preference (no Env/DEFAULT fallback).

CORE_TOOLS
DEFAULT_PREREQUISITES = 
{ 'shell' => %w[shell].freeze, 'pwn_eval' => %w[ruby].freeze }.freeze
DOMAIN_PINS = 
{
  re: %w[pty_open pty_send pty_read pty_close binary_triage exploitdev job_run decompile],
  pentest: %w[pty_open job_run finding_record binary_triage fuzz_campaign exploitdev],
  fuzzing: %w[fuzz_campaign job_run job_status job_result],
  exploitation: %w[exploitdev pty_open job_run finding_record],
  reporting: %w[finding_record finding_report artifacts_list]
}.freeze

Class Method Summary collapse

Class Method Details

.all ⇒ Object

Supported Method Parameters

entries = PWN::AI::Agent::Registry.all



99
100
101
 
# File 'lib/pwn/ai/agent/registry.rb', line 99

public_class_method def self.all
  @entries.values
end

.apply_preference(opts = {}) ⇒ Object

Supported Method Parameters

sorted = PWN::AI::Agent::Registry.apply_preference( items: 'required - Array of Entry objects or OpenAI tool hashes', order: 'optional - forwarded to preference_order only when the key is set', preference: 'optional - alias of :order; same key-present rule' )

Stable sort: preferred names first, original index as the tie-break so equal-preference items keep their incoming order. Empty preference (explicit [] / nil) leaves items unsorted.



181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
 
# File 'lib/pwn/ai/agent/registry.rb', line 181

public_class_method def self.apply_preference(opts = {})
  items = Array(opts[:items] || opts[:entries])
  po = {}
  po[:preference] = opts[:preference] if opts.key?(:preference)
  po[:order] = opts[:order] if opts.key?(:order)
  order = preference_order(po)
  return items if order.empty?

  index = {}
  order.each_with_index { |name, i| index[name.to_s] = i }
  fallback = order.length
  items.sort_by.with_index do |item, orig|
    name = if item.respond_to?(:name)
             item.name.to_s
           elsif item.is_a?(Hash)
             (item.dig(:function, :name) ||
              item.dig('function', 'name') ||
              item[:name] || item['name']).to_s
           else
             item.to_s
           end
    [index.fetch(name, fallback), orig]
  end
end

.authors ⇒ Object

Author(s)

0day Inc. [email protected]



413
414
415
 
# File 'lib/pwn/ai/agent/registry.rb', line 413

public_class_method def self.authors
  "AUTHOR(S):\n  0day Inc. <[email protected]>\n"
end

.available?(opts = {}) ⇒ Boolean

Unknown capability is not an absence. Only trusted negative evidence excludes a tool; learned history and CORE/domain pins cannot override it.

Returns:

  • (Boolean) 


351
352
353
354
355
356
357
358
359
360
361
362
 
# File 'lib/pwn/ai/agent/registry.rb', line 351

public_class_method def self.available?(opts = {})
  entry = opts[:entry] || lookup(name: opts[:name])
  return true unless entry && defined?(Policy)

  raw = opts[:trusted_context] || Policy.current_episode&.dig(:trusted_context)
  return true unless raw.is_a?(Hash)

  observed = Policy.observed_context(trusted_context: raw)
  Array(entry.prerequisites).none? do |capability|
    observed[:missing_prerequisites].include?(capability.to_s) || observed[:capabilities][capability.to_sym] == false
  end
end

.definitions(opts = {}) ⇒ Object

Supported Method Parameters

tools = PWN::AI::Agent::Registry.definitions( enabled: 'optional - Array of toolset names to include; nil = all whose check passes', relevance: 'optional - user request; when set AND :tool_router is enabled, slim to CORE + top-K keyword matches', top_k: 'optional - keyword-ranked tools to keep beyond CORE (default 10)', order: 'optional - preference list forwarded only when this key is present', preference: 'optional - alias of :order; same key-present rule' )



119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
 
# File 'lib/pwn/ai/agent/registry.rb', line 119

public_class_method def self.definitions(opts = {})
  eager_load!
  enabled = opts[:enabled]
  enabled = enabled.map(&:to_s) if enabled
  pool = @entries.values.select { |e| (enabled.nil? || enabled.include?(e.toolset)) && safe_check(entry: e) }
  pool = pool.select { |e| available?(entry: e, trusted_context: opts[:trusted_context]) }

  pref_fwd = {}
  pref_fwd[:order] = opts[:order] if opts.key?(:order)
  pref_fwd[:preference] = opts[:preference] if opts.key?(:preference)
  pref_fwd[:kind] = opts[:kind] if opts.key?(:kind)
  pref_fwd[:intent] = opts[:intent] if opts.key?(:intent)
  pref_fwd[:trusted_context] = opts[:trusted_context] if opts.key?(:trusted_context)

  if opts[:core_only]
    pool = pool.select { |e| CORE_TOOLS.include?(e.name) }
  elsif opts[:relevance] && router_enabled?
    keep = rank({ query: opts[:relevance], entries: pool }.merge(pref_fwd)).first(opts[:top_k] || 10).map(&:name)
    pinned = domain_pin_names(query: opts[:relevance])
    names = (CORE_TOOLS + keep + pinned).uniq
    pool  = pool.select { |e| names.include?(e.name) }
  end

  pool = apply_preference({ items: pool }.merge(pref_fwd))
  pool.map { |e| { type: 'function', function: e.schema } }
end

.discover(opts = {}) ⇒ Object

Supported Method Parameters

names = PWN::AI::Agent::Registry.discover( force: 'optional - re-require tool files even if already discovered (default false)' )



314
315
316
 
# File 'lib/pwn/ai/agent/registry.rb', line 314

public_class_method def self.discover(opts = {})
  eager_load!(opts)
end

.eager_load!(opts = {}) ⇒ Object 



318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
 
# File 'lib/pwn/ai/agent/registry.rb', line 318

public_class_method def self.eager_load!(opts = {})
  force = opts[:force] ||= false
  return @entries.keys if @discovered && !force

  tools_dir = File.join(__dir__, 'tools')
  if Dir.exist?(tools_dir)
    Dir[File.join(tools_dir, '*.rb')].each do |f|
      require f
    rescue StandardError, LoadError => e
      warn "[pwn-ai] failed to load tool #{File.basename(f)}: #{e.class}: #{e.message}"
    end
  end
  require 'pwn/ai/agent/manifest'
  Manifest.register
  @discovered = true
  @entries.keys
end

.help ⇒ Object

Display Usage for this Module



419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
 
# File 'lib/pwn/ai/agent/registry.rb', line 419

public_class_method def self.help
  puts "USAGE:
    # Run register and return its result
    #{self}.register(
      name: 'required - tool name exposed to the model',
      toolset: 'required - grouping for enable/disable (terminal, file, pwn, memory…)',
      schema: 'required - OpenAI function schema {name:, description:, parameters:}',
      handler: 'required - ->(args_hash) { ... } returning a JSON-serialisable object',
      check: 'optional - -> { bool } gate; tool only advertised when truthy',
      prerequisites: 'optional - required capability categories; observed absence excludes this tool',
      max_chars: 'optional - cap on serialised result (default 24_000)'
    )

    # Run lookup and return its result
    #{self}.lookup(
      name: 'required - registered tool name'
    )

    # Run all and return its result
    #{self}.all

    # Run toolsets and return its result
    #{self}.toolsets

    # Run definitions and return its result
    #{self}.definitions(
      enabled: 'optional - Array of toolset names to include; nil = all whose check passes',
      relevance: 'optional - user request; when set AND :tool_router is enabled, slim to CORE + top-K keyword matches',
      top_k: 'optional - keyword-ranked tools to keep beyond CORE (default 10)',
      order: 'optional - preference list forwarded only when this key is present',
      preference: 'optional - alias of :order; same key-present rule',
      kind: 'optional - kind value consumed by #definitions',
      intent: 'optional - intent value consumed by #definitions',
      core_only: 'optional - core only value consumed by #definitions',
      trusted_context: 'optional - caller observations accepted by Policy.observed_context; defaults to live episode observations'
    )

    # Run preference order and return its result
    #{self}.preference_order(
      order: 'optional - Array of tool names; explicit nil/[] disables',
      preference: 'optional - alias of :order. Explicit key disables Env/DEFAULT fallback'
    )

    # Run apply preference and return its result
    #{self}.apply_preference(
      items: 'required - Array of Entry objects or OpenAI tool hashes',
      order: 'optional - forwarded to preference_order only when the key is set',
      preference: 'optional - alias of :order; same key-present rule',
      entries: 'optional - entries value consumed by #apply_preference'
    )

    # Run rank and return its result
    #{self}.rank(
      query: 'required - user request text',
      entries: 'optional - Entry pool to rank (default .all)',
      order: 'optional - preference list forwarded to preference_order',
      preference: 'optional - preference value consumed by #rank',
      trusted_context: 'optional - caller-only environment and prerequisite observations; missing tools cannot win via history'
    )

    # Test observed prerequisites independently of learned ranking.
    #{self}.available?(
      entry: 'optional - registered Entry',
      name: 'optional - lookup name when entry is absent',
      trusted_context: 'optional - observed capabilities and missing prerequisites; default live episode context'
    )

    # Run discover and return its result
    #{self}.discover(
      force: 'optional - re-require tool files even if already discovered (default false)'
    )

    # Eager-load every tools/*.rb so Registry.definitions works in a cold REPL.
    #{self}.eager_load!(
      force: 'optional - re-require tool files even if already discovered (default false)'
    )

    # Schema/handler contract check for every registered tool.
    #{self}.selftest(
      dry_run: 'optional - reserved dry-run flag'
    )

    # Print the AUTHOR(S) string for this module.
    #{self}.authors
  "
  constants.sort
end

.lookup(opts = {}) ⇒ Object

Supported Method Parameters

entry = PWN::AI::Agent::Registry.lookup( name: 'required - registered tool name' )



91
92
93
94
 
# File 'lib/pwn/ai/agent/registry.rb', line 91

public_class_method def self.lookup(opts = {})
  name = opts[:name]
  @entries[name.to_s]
end

.preference_order(opts = {}) ⇒ Object

Supported Method Parameters

order = PWN::AI::Agent::Registry.preference_order( order: 'optional - Array of tool names; explicit nil/[] disables', preference: 'optional - alias of :order. Explicit key disables Env/DEFAULT fallback' )



152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
 
# File 'lib/pwn/ai/agent/registry.rb', line 152

public_class_method def self.preference_order(opts = {})
  explicit = opts.key?(:order) || opts.key?(:preference)
  if explicit
    raw = opts.key?(:preference) ? opts[:preference] : opts[:order]
    return [] if raw.nil? || (raw.respond_to?(:empty?) && raw.empty?)

    return Array(raw).map(&:to_s).reject(&:empty?)
  end

  raw = nil
  raw = PWN::Env.dig(:ai, :agent, :tool_preference) if defined?(PWN::Env) && PWN::Env.is_a?(Hash)
  return Array(raw).map(&:to_s).reject(&:empty?) unless raw.nil? || (raw.respond_to?(:empty?) && raw.empty?)

  DEFAULT_PREFERENCE.dup
rescue StandardError
  DEFAULT_PREFERENCE.dup
end

.rank(opts = {}) ⇒ Object

Supported Method Parameters

ranked = PWN::AI::Agent::Registry.rank( query: 'required - user request text', entries: 'optional - Entry pool to rank (default .all)', order: 'optional - preference list forwarded to preference_order' )

C1 advantage-weighted router: score = α·keyword_sim + β·(tool rolling success_rate − global rate) + γ·UCB1(tool). Untried tools get an exploration bonus; tools that outperform the fleet get an exploitation bonus. Thompson sampling is available via Metrics.thompson for stochastic routing.



219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
 
# File 'lib/pwn/ai/agent/registry.rb', line 219

public_class_method def self.rank(opts = {})
  query   = opts[:query].to_s.downcase
  entries = opts[:entries] || all
  entries = entries.select { |entry| available?(entry: entry, trusted_context: opts[:trusted_context]) }
  return entries if query.strip.empty?

  tokens = query.scan(/[a-z0-9_]{3,}/).uniq
  # C1 — advantage-weighted router:
  #   score = alpha*keyword_sim + beta*advantage + gamma*UCB(tool)
  #         + delta*prm_advantage + eps_q*Q_adv + zeta*pref_score
  # UCB gives untried / low-N tools an exploration bonus so a single
  # early failure (before its dep was installed) does not blacklist
  # it forever; advantage prefers tools that outperform the fleet.
  # zeta * pref_score (zeta=0.35) boosts Env/opts tool preference;
  # alpha stays 1.0 so keyword fit remains the primary signal.
  # zeta is 0 when preference is empty so ties are not pref-boosted.
  alpha = 1.0
  po = {}
  po[:preference] = opts[:preference] if opts.key?(:preference)
  po[:order] = opts[:order] if opts.key?(:order)
  pref_order = preference_order(po)
  zeta = pref_order.empty? ? 0.0 : 0.35
  pref_index = {}
  pref_order.each_with_index { |n, i| pref_index[n.to_s] = i }
  pref_n = [pref_order.length, 1].max
  # P4 — haircut advantage weight when reward proxy is hacked
  trust = defined?(Metrics) && Metrics.respond_to?(:proxy_trust) ? Metrics.proxy_trust : 1.0
  beta  = 0.3 * trust
  gamma = 0.2
  # P18/P2 — PRM step_reward advantage closes R2 into the controller.
  # Sample-efficiency: if fewer than 3 tools have prm_n≥PRM_MIN_N,
  # drop delta to 0 so sparse PRM cannot inject rank variance.
  # Otherwise scale delta by fleet coverage fraction.
  prm_ready = 0
  if defined?(Metrics) && Metrics.respond_to?(:prm_n)
    prm_ready = entries.count do |e|
      Metrics.prm_n(name: e.name).to_i >= begin
        Metrics::PRM_MIN_N
      rescue StandardError
        5
      end
    end
  end
  fleet = [entries.length, 1].max
  coverage = prm_ready.to_f / fleet
  delta = if prm_ready < 3
            0.0
          else
            0.25 * trust * [coverage / 0.3, 1.0].min
          end
  # R5 — tabular Q(s,a)−V(s). Weight sits next to Metrics advantage
  # so a warm table actually moves rank. Keyword fit stays first
  # (alpha=1.0); CORE_TOOLS still always included.
  eps_q = if defined?(PWN::AI::Agent::Policy) && Policy.respond_to?(:advantage) && Policy.respond_to?(:enabled?) && Policy.enabled?
            warm = Policy.respond_to?(:episode_budget_met?) && Policy.episode_budget_met?
            (warm ? 0.45 : 0.20) * [trust, 0.55].max
          else
            0.0
          end
  pol_state = (Policy.current_state || Policy.state(request: query) if defined?(PWN::AI::Agent::Policy) && Policy.respond_to?(:current_state))
  pol_state = Policy.observed_state(state: pol_state, trusted_context: opts[:trusted_context]) if defined?(Policy) && opts.key?(:trusted_context)
  # Metrics has no environment dimension. Do not transfer its global
  # success/PRM/UCB history into explicitly observed environments.
  if pol_state.to_s.include?('|env:')
    beta = 0.0
    gamma = 0.0
    delta = 0.0
  end
  scored = entries.map do |e|
    hay   = "#{e.name} #{e.toolset} #{e.schema[:description]} #{Array(e.schema.dig(:parameters, :properties)&.keys).join(' ')}".downcase
    sim   = tokens.count { |t| hay.include?(t) }
    adv   = defined?(Metrics) && Metrics.respond_to?(:advantage) ? Metrics.advantage(name: e.name) : 0.0
    ucb   = defined?(Metrics) && Metrics.respond_to?(:ucb) ? Metrics.ucb(name: e.name) : 0.5
    prm   = defined?(Metrics) && Metrics.respond_to?(:prm_advantage) ? Metrics.prm_advantage(name: e.name) : 0.0
    qadv  = if eps_q.positive? && pol_state
              Policy.advantage(state: pol_state, action: e.name, context_state: Policy.current_context_state)
            else
              0.0
            end
    pidx = pref_index[e.name]
    pref = pidx.nil? ? 0.0 : ((pref_n - pidx).to_f / pref_n)
    total = (alpha * sim) + (beta * adv) + (gamma * ucb) + (delta * prm) + (eps_q * qadv) + (zeta * pref)
    [e, sim, total]
  end
  # Equal keyword-fit (and equal total): preferred name first.
  scored.reject { |_, sim, _| sim.zero? }
        .sort_by { |e, _, s| [-s, pref_index.fetch(e.name, pref_n), e.name] }
        .map(&:first)
end

.register(opts = {}) ⇒ Object

Supported Method Parameters

PWN::AI::Agent::Registry.register( name: 'required - tool name exposed to the model', toolset: 'required - grouping for enable/disable (terminal, file, pwn, memory…)', schema: 'required - OpenAI function schema description:, parameters:', handler: 'required - ->(args_hash) { ... } returning a JSON-serialisable object', check: 'optional - -> { bool } gate; tool only advertised when truthy', prerequisites: 'optional - required capability categories; shell and pwn_eval have core defaults', max_chars: 'optional - cap on serialised result (default 24_000)' )



69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
 
# File 'lib/pwn/ai/agent/registry.rb', line 69

public_class_method def self.register(opts = {})
  name = opts[:name].to_s
  raise 'ERROR: name is required' if name.empty?
  raise 'ERROR: schema is required' unless opts[:schema]
  raise 'ERROR: handler is required' unless opts[:handler].respond_to?(:call)

  @entries[name] = Entry.new(
    name: name,
    toolset: opts[:toolset].to_s,
    schema: opts[:schema],
    handler: opts[:handler],
    check: opts[:check] ||= -> { true },
    max_chars: opts[:max_chars] ||= 24_000,
    prerequisites: Array(opts.fetch(:prerequisites, DEFAULT_PREREQUISITES[name])).map(&:to_s)
  )
end

.selftest(opts = {}) ⇒ Object 



336
337
338
339
340
341
342
343
344
345
346
347
 
# File 'lib/pwn/ai/agent/registry.rb', line 336

public_class_method def self.selftest(opts = {})
  _dry = opts[:dry_run]
  eager_load!
  tools_dir = File.join(__dir__, 'tools')
  files = Dir[File.join(tools_dir, '*.rb')]
  rows = @entries.map do |name, e|
    schema = e.schema.is_a?(Hash)
    required = Array(e.schema&.dig(:parameters, :required) || e.schema&.dig('parameters', 'required'))
    { name: name, schema_valid: schema, handler: !e.handler.nil?, required: required }
  end
  { tools: rows, files: files.length, registered: @entries.length, ok: rows.length >= files.length }
end

.toolsets ⇒ Object

Supported Method Parameters

names = PWN::AI::Agent::Registry.toolsets



106
107
108
 
# File 'lib/pwn/ai/agent/registry.rb', line 106

public_class_method def self.toolsets
  @entries.values.map(&:toolset).uniq.sort
end