Module: PWN::AI::Agent::Dispatch

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

Overview

Tool-call dispatch: takes a single tool_call object (OpenAI shape), looks up the registered handler, parses args, runs it, and returns a JSON string suitable for a role:'tool' message.

TOLERANT DISPATCH (local-model scaffolding)

Local models running on Ollama frequently emit almost- right tool calls: run_shell instead of shell, trailing commas, single-quoted JSON, arguments as a bare string. Strict parsing burns an iteration and often spirals. Dispatch now:

* repair_name  — Levenshtein-matches unknown names to the closest
               registered tool and records a Mistakes fingerprint
               (source: :repair) so the KNOWN MISTAKES block
               eventually teaches the model the right name.
* parse_args   — falls back to a JSON5-ish clean-up pass (strip
               trailing commas, swap single→double quotes, wrap a
               bare scalar as the tool's sole required arg).

Frontier engines never hit these paths — repair is a no-op when the name/JSON are already valid.

Constant Summary collapse

WRITE_ARGV_RX =

Effect of a tool call from NAME + ARGV only — never stdout. :write mutate, :browse navigate, :recall store lookup, :read/:eval observe.

/
  \bsed\s+-i\b|\bruby\s+-i\b|\btee\b|
  (?:\s|\A)>{1,2}\s+\S|
  File\.(?:write|open|binwrite)|IO\.write|
  \bopen\s*\([^)]*['"]w|
  \b(?:cp|mv|rm|mkdir|touch|chmod|chown)\b|
  \bgit\s+(?:add|commit|rm)
/ix
BROWSE_ARGV_RX = 
/
  TransparentBrowser|browser_obj|\.goto\b|dump_links|
  watir|headless_?chrome|\bdevtools\b
/ix
RECALL_TOOLS = 
%w[
  memory_recall session_recall skills_recall sessions_view
  sessions_list sessions_current
].freeze
STORE_TOOLS = 
%w[
  memory_remember mistakes_record mistakes_resolve
  learning_note_outcome skill_create skill_add_reference skills_update
].freeze

Class Method Summary collapse

Class Method Details

.authors ⇒ Object

Author(s)

0day Inc. [email protected]



598
599
600
 
# File 'lib/pwn/ai/agent/dispatch.rb', line 598

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

.call(opts = {}) ⇒ Object

Supported Method Parameters

json_str = PWN::AI::Agent::Dispatch.call( tool_call: 'required - Hash { id:, type:, function: { name:, arguments: } }' )



39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
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
145
146
147
148
149
150
 
# File 'lib/pwn/ai/agent/dispatch.rb', line 39

public_class_method def self.call(opts = {})
  tool_call = opts[:tool_call]
  raise 'ERROR: tool_call is required' if tool_call.nil?

  fn   = tool_call[:function] || tool_call['function'] || {}
  name = (fn[:name] || fn['name']).to_s
  raw  = if fn.key?(:arguments) || fn.key?('arguments')
           fn[:arguments] || fn['arguments']
         else
           '{}'
         end

  entry = Registry.lookup(name: name) || Registry.lookup(name: repair_name(name: name))
  return JSON.generate(error: "unknown tool: #{name}") unless entry

  args = parse_args(raw: raw, entry: entry)
  args = alias_known_keys(args: args, entry: entry)
  schema = entry.schema[:parameters] || entry.schema['parameters'] || { type: 'object' }
  declaration = Manifest.load(directory: opts[:manifest_directory] || Manifest::DIRECTORY)[entry.name]
  schema = { allOf: [schema, declaration['params']] } if declaration
  if raw.nil?
    required = Array(schema[:required] || schema['required']).map(&:to_s)
    missing = required.reject { |key| ToolGuard.present?(value: args[key.to_sym] || args[key]) }
    return schema_denial(type: 'required', error: "Supply required fields: #{missing.join(', ')}") unless missing.empty?
  end
  type_schema = drop_required(node: schema)
  violation = JSONSchemer.schema(JSON.parse(JSON.generate(type_schema))).validate(JSON.parse(JSON.generate(args))).first
  return schema_denial(violation.transform_keys(&:to_sym)) if violation

  blob = args.inspect
  if defined?(PWN::Plugins::Vault)
    blob = PWN::Plugins::Vault.expand(text: blob)
    args = expand_vault_args(args: args)
  end
  denied = Manifest.check(opts.merge(name: entry.name, args: args))
  return JSON.generate(denied) if denied

  if defined?(Engagement)
    denied = ToolGuard.scope_check!(args: args, command: blob) if defined?(ToolGuard) && ToolGuard.respond_to?(:scope_check!)
    denied ||= Engagement.deny_if_out_of_scope(args: args, command: blob)
    return JSON.generate(denied) if denied
  end
  if defined?(ToolGuard) && ToolGuard.respond_to?(:policy_decision)
    pol = ToolGuard.policy_decision(name: entry.name, args: args)
    return JSON.generate(pol) if pol.is_a?(Hash) && pol[:action] == 'deny'
  end
  return JSON.generate(success: false, error: 'taint: tool-output instruction in args', code: 'TAINT_DENY') if taint_blocked?(name: entry.name, args: args)
  if defined?(ToolGuard) && ToolGuard.respond_to?(:canary_leak?) &&
     ToolGuard.canary_leak?(text: args.inspect)
    return JSON.generate(success: false, error: 'refused: session canary in outbound args', code: 'CANARY_DENY', rule_id: 'canary')
  end
  if defined?(ToolGuard) && ToolGuard.respond_to?(:refuse_copied_persist?) &&
     ToolGuard.refuse_copied_persist?(name: entry.name, args: args)
    return JSON.generate(
      success: false,
      error: 'refused: memory_remember/skills_update text copied from last tool output'
    )
  end
  if blob.match?(/open_sockraw|sockraw/) && defined?(PWN::Plugins::PreflightChecker) &&
     PWN::Plugins::PreflightChecker.respond_to?(:cap_net_raw?) && !PWN::Plugins::PreflightChecker.cap_net_raw?
    return JSON.generate(
      success: false,
      error: 'capability missing CAP_NET_RAW',
      substitute: 'PWN::Plugins::Packet.tcp_connect_scan',
      code: 'CAP_DENY'
    )
  end
  mission = defined?(Mission) ? Mission.active : nil
  ack = Confirmation.gate(
    name: entry.name,
    args: args,
    engagement_id: opts[:engagement_id] || args[:engagement_id] || args['engagement_id'],
    operator_ack: opts[:operator_ack] || args[:operator_ack] || args['operator_ack'],
    scope_path: opts[:scope_path],
    unattended: mission && mission[:unattended] == true,
    ack_scope: mission && mission[:unattended] == true
  )
  return JSON.generate(ack) if ack

  budget = prepare_budget(opts.merge(entry: entry, args: args))
  return JSON.generate(budget[:denial]) if budget && budget[:denial]

  started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
  begin
    result = entry.handler.call(args)
  rescue StandardError => e
    finish_budget(budget: budget, result: { error: e.is_a?(Timeout::Error) ? 'timeout' : e.class.name }, elapsed: Process.clock_gettime(Process::CLOCK_MONOTONIC) - started)
    raise
  end
  telemetry = finish_budget(budget: budget, result: result, elapsed: Process.clock_gettime(Process::CLOCK_MONOTONIC) - started)
  result = Result.page(value: result, entry: entry, session_id: opts[:session_id])
  result = ToolGuard.quarantine_output(text: result) if defined?(ToolGuard) && result.is_a?(String) && ToolGuard.respond_to?(:quarantine_output)
  note_taint(text: result)
  if defined?(PWN::Plugins::Vault) && result.is_a?(String)
    result = PWN::Plugins::Vault.redact(text: result)
  elsif defined?(PWN::Plugins::Vault) && result.is_a?(Hash)
    result = JSON.parse(PWN::Plugins::Vault.redact(text: JSON.generate(result)))
  end
  response = { success: true, result: result, effect: effect(name: entry.name, args: args) }
  response[:budget] = telemetry if telemetry
  if telemetry && telemetry[:remaining_s].zero?
    response[:success] = false
    response[:error] = 'budget_exhausted'
  end
  JSON.generate(response)
rescue StandardError => e
  JSON.generate(
    success: false,
    error: "#{e.class}: #{e.message}",
    backtrace: Array(e.backtrace).first(3)
  )
end

.effect(opts = {}) ⇒ Object 



514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
 
# File 'lib/pwn/ai/agent/dispatch.rb', line 514

public_class_method def self.effect(opts = {})
  name = opts[:name].to_s
  return :read if name.empty?
  return :recall if RECALL_TOOLS.include?(name)
  return :store if STORE_TOOLS.include?(name)

  blob = argv_blob(args: opts[:args])
  return :write if blob.match?(WRITE_ARGV_RX)
  return :browse if blob.match?(BROWSE_ARGV_RX)
  return :eval if name == 'pwn_eval'

  :read
rescue StandardError
  :read
end

.help ⇒ Object

Display Usage for this Module



604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
 
# File 'lib/pwn/ai/agent/dispatch.rb', line 604

public_class_method def self.help
  puts "USAGE:
    # Run call and return its result
    #{self}.call(
      tool_call: 'required - Hash { id:, type:, function: { name:, arguments: } }',
      manifest_directory: 'optional - trusted manifest YAML directory',
      budget_ledger: 'optional - caller-owned Hash retained across one task',
      budget_key: 'optional - trusted approach identifier; defaults to tool name',
      scope_policy: 'optional - trusted policy Hash',
      scope_path: 'optional - trusted scope file path',
      session_id: 'optional - session directory for oversized result artifacts',
      audit_path: 'optional - trusted audit JSONL path',
      approval_callback: 'optional - trusted callback for prompt risk gates',
      engagement_id: 'optional - engagement id used to cache exploit/destructive ACK',
      operator_ack: 'optional - true records a one-time ACK for this engagement'
    )

    # Run repair name and return its result
    #{self}.repair_name(
      name: 'required - possibly-wrong tool name emitted by the model'
    )

    # Run tool calls from text and return its result
    #{self}.tool_calls_from_text(
      text: 'required - assistant plain-text that may embed shell(...) / JSON tool forms',
      call: 'optional - shell{command: uname -s} / tool:shell{command:id}'
    )

    # Run effect and return its result
    #{self}.effect(
      name: 'required - binary or identifier name',
      args: 'optional - args value consumed by #effect'
    )

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

.repair_name(opts = {}) ⇒ Object

Supported Method Parameters

fixed = PWN::AI::Agent::Dispatch.repair_name( name: 'required - possibly-wrong tool name emitted by the model' )

Returns the closest registered tool name by Levenshtein distance (max distance = 1/3 of the emitted name, min 3) or nil when nothing is close enough. Every successful repair is fingerprinted into Mistakes so the negative-feedback loop trains the model's output format via its own system prompt.



163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
 
# File 'lib/pwn/ai/agent/dispatch.rb', line 163

public_class_method def self.repair_name(opts = {})
  name = opts[:name].to_s
  return nil if name.empty?

  pool = Registry.all.map(&:name)
  return nil if pool.empty?

  best, dist = pool.map { |n| [n, DidYouMean::Levenshtein.distance(n, name)] }
                   .min_by(&:last)
  thresh = [(name.length / 3.0).ceil, 3].max
  return nil if dist > thresh

  if defined?(Mistakes)
    Mistakes.record(
      tool: 'tool_name',
      error: "model emitted '#{name}', repaired to '#{best}'",
      args: name,
      source: :repair
    )
  end
  best
rescue StandardError
  nil
end

.tool_calls_from_text(opts = {}) ⇒ Object

Supported Method Parameters

calls = PWN::AI::Agent::Dispatch.tool_calls_from_text( text: 'required - assistant plain-text that may embed shell(...) / JSON tool forms' )

Local / abliterated models often print tool invocations as content instead of native message.tool_calls. Supported shapes include:

shell(command="id") / shell({"command":"id"}) / shell("id")
{"name":"shell","arguments":{...}} / {"function":{"name":...}}
{"tool":"shell","arguments":{...}} / {"call":"shell","arguments":{...}}
call:shell{command: "uname -s"} / tool:shell{"command":"id"}

When structured tool_calls are empty, Loop coerces those strings into OpenAI-shaped tool_call hashes so Dispatch runs them instead of treating the string as a FINAL answer.



347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
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
 
# File 'lib/pwn/ai/agent/dispatch.rb', line 347

public_class_method def self.tool_calls_from_text(opts = {})
  text = opts[:text].to_s
  return [] if text.strip.empty?

  Registry.discover if defined?(Registry) && Registry.respond_to?(:discover)
  known = if defined?(Registry)
            Registry.all.map { |e| e.name.to_s }.reject(&:empty?)
          else
            %w[shell pwn_eval]
          end
  return [] if known.empty?

  names_alt = known.map { |n| Regexp.escape(n) }.join('|')
  calls = []
  seen = {}

  add = lambda do |name, args|
    name = name.to_s
    next unless known.include?(name)

    args_h = case args
             when Hash then symbolize(hash: args)
             when String
               s = args.strip
               begin
                 parsed = JSON.parse(s, symbolize_names: true)
                 parsed.is_a?(Hash) ? parsed : { value: parsed }
               rescue JSON::ParserError
                 h = {}
                 s.scan(/([A-Za-z_]\w*)\s*[:=]\s*(?:"((?:\\.|[^"])*)"|'((?:\\.|[^'])*)'|([^\s,)}{]+))/) do
                   k = Regexp.last_match(1)
                   h[k.to_sym] = Regexp.last_match(2) || Regexp.last_match(3) || Regexp.last_match(4)
                 end
                 if h.empty?
                   entry = (Registry.lookup(name: name) if defined?(Registry))
                   req = Array(entry&.schema&.dig(:parameters, :required))
                   h = req.length == 1 ? { req.first.to_sym => s } : { command: s }
                 end
                 h
               end
             else
               {}
             end
    key = "#{name}|#{JSON.generate(args_h)}"
    next if seen[key]

    seen[key] = true
    calls << {
      id: "textcall_#{calls.length + 1}_#{SecureRandom.hex(3)}",
      type: 'function',
      function: {
        name: name,
        # OpenAI/xAI wire format requires a JSON string, not a map.
        arguments: JSON.generate(args_h)
      }
    }
  end

  # Balanced-delimiter extractor used for name(...) and call:name{...}.
  extract_balanced = lambda do |open_ch, close_ch, from|
    depth = 1
    i = from
    in_s = nil
    esc = false
    while i < text.length && depth.positive?
      ch = text[i]
      if in_s
        if esc
          esc = false
        elsif ch == '\\'
          esc = true
        elsif ch == in_s
          in_s = nil
        end
      elsif ['"', "'"].include?(ch)
        in_s = ch
      elsif ch == open_ch
        depth += 1
      elsif ch == close_ch
        depth -= 1
      end
      i += 1
    end
    depth.zero? ? [text[from...(i - 1)].to_s.strip, i] : nil
  end

  # JSON object forms:
  #   {"name":"shell","arguments":{...}}
  #   {"function":{"name":"shell","arguments":{...}}}
  #   {"tool":"shell","arguments":{...}} / {"call":"shell",...}
  #   {"type":"call","name":"shell",...}
  text.scan(/\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}/m).each do |blob|
    begin
      j = JSON.parse(blob, symbolize_names: true)
    rescue JSON::ParserError
      next
    end
    next unless j.is_a?(Hash)

    name = (
      j[:name] || j[:tool] || j[:call] ||
      j.dig(:function, :name) || j.dig(:tool_call, :name)
    ).to_s
    # Skip pure type tags mistaken as names (e.g. {"call":{...}} trees).
    next if name.empty? || %w[function tool_call].include?(name)

    args = j[:arguments] || j[:args] || j[:parameters] ||
           j.dig(:function, :arguments) || j.dig(:tool_call, :arguments) || {}
    add.call(name, args)
  end

  # Colon-brace forms (OpenWebUI / abliterated dumps):
  #   call:shell{command: "uname -s"}
  #   tool:shell{"command":"id"}
  #   call:shell{command="id"}
  rx_colon = /\b(?:call|tool)\s*:\s*(#{names_alt})\s*\{/i
  idx = 0
  while (m = text.match(rx_colon, idx))
    name = m[1]
    extracted = extract_balanced.call('{', '}', m.end(0))
    if extracted
      # Re-wrap: balanced extractor yields the interior only. Paren form
      # shell({...}) keeps braces inside (...); brace form must restore
      # them so JSON.parse / kwarg scan see a full object body.
      add.call(name, "{#{extracted[0]}}")
    end
    idx = m.begin(0) + 1
  end

  # Call forms: shell(command="...") / shell({"command":"id"}) / shell("id")
  rx = /\b(#{names_alt})\s*\(/i
  idx = 0
  while (m = text.match(rx, idx))
    name = m[1]
    extracted = extract_balanced.call('(', ')', m.end(0))
    add.call(name, extracted[0]) if extracted
    idx = m.begin(0) + 1
  end

  calls
rescue StandardError
  []
end