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.

Class Method Summary collapse

Class Method Details

.authors ⇒ Object

Author(s)

0day Inc. support@0dayinc.com



288
289
290
 
# File 'lib/pwn/ai/agent/dispatch.rb', line 288

public_class_method def self.authors
  "AUTHOR(S):\n  0day Inc. <support@0dayinc.com>\n"
end

.call(opts = {}) ⇒ Object

Supported Method Parameters

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



34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
 
# File 'lib/pwn/ai/agent/dispatch.rb', line 34

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  = fn[:arguments] || fn['arguments'] || '{}'

  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)
  result = entry.handler.call(args)
  JSON.generate(success: true, result: result)
rescue StandardError => e
  JSON.generate(
    success: false,
    error: "#{e.class}: #{e.message}",
    backtrace: Array(e.backtrace).first(3)
  )
end

.help ⇒ Object

Display Usage for this Module



294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
 
# File 'lib/pwn/ai/agent/dispatch.rb', line 294

public_class_method def self.help
  puts <<~USAGE
    USAGE:
      json_str = PWN::AI::Agent::Dispatch.call(
        tool_call: {
          id: 'call_1',
          type: 'function',
          function: { name: 'shell', arguments: '{"command":"id"}' }
        }
      )

      PWN::AI::Agent::Dispatch.repair_name(name: 'run_shell')  # => 'shell'
      PWN::AI::Agent::Dispatch.tool_calls_from_text(text: 'shell(command="id")')

      #{self}.authors
  USAGE
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.



67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
 
# File 'lib/pwn/ai/agent/dispatch.rb', line 67

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.



137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
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
 
# File 'lib/pwn/ai/agent/dispatch.rb', line 137

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