Module: PWN::AI::Agent::TaskDAG

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

Overview

Planner/executor split: YAML task DAGs with per-step checkpoints.

Constant Summary collapse

ROOT = 
File.join(Dir.home, '.pwn', 'runs')

Class Method Summary collapse

Class Method Details

.authors ⇒ Object 



70
71
72
 
# File 'lib/pwn/ai/agent/task_dag.rb', line 70

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

.execute(opts = {}) ⇒ Object

Run an approved DAG; completed checkpoints are skipped on resume.



35
36
37
38
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
 
# File 'lib/pwn/ai/agent/task_dag.rb', line 35

public_class_method def self.execute(opts = {})
  refuse_operator_shell!(opts)
  root = opts[:root] || ROOT
  run_id = (opts[:run_id] || SecureRandom.hex(8)).to_s
  dir = File.join(root, run_id)
  FileUtils.mkdir_p(dir)
  dag = load_dag(opts.merge(dir: dir))
  refuse_unapproved!(dag: dag, opts: opts)
  File.write(File.join(dir, 'dag.yaml'), YAML.dump(JSON.parse(JSON.generate(dag)))) unless File.file?(File.join(dir, 'dag.yaml'))
  File.write(File.join(dir, 'run_id'), run_id)
  done = load_done(dir: dir)
  steps = Array(dag['steps'] || dag[:steps]).map { |step| normalize_step(step: step) }
  results = []
  remaining = steps.reject { |step| done[step['id']] }
  until remaining.empty?
    ready = remaining.select { |step| Array(step['dependencies']).all? { |need| done[need] } }
    raise ArgumentError, 'task DAG deadlock or missing dependencies' if ready.empty?

    ready.each do |step|
      remaining.delete(step)
      row = run_step(step: step, dir: dir)
      done[step['id']] = true
      results << row
      persist_checkpoint(dir: dir, step: step, row: row)
      return { run_id: run_id, dir: dir, results: results, halted: step['id'] } if opts[:halt_after].to_s == step['id']
    end
  end
  { run_id: run_id, dir: dir, results: results, ok: results.all? { |row| row[:ok] != false } }
end

.help ⇒ 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
103
104
105
106
107
108
 
# File 'lib/pwn/ai/agent/task_dag.rb', line 74

public_class_method def self.help
  puts "USAGE:
    # List host binaries this module expects to be installed.
    #{self}.required_bins

    # Emit a YAML-serializable DAG with zero tool side effects.
    #{self}.plan(
      request: 'required - operator request to decompose',
      steps: 'optional - Array of step hashes; inferred from the request when omitted'
    )

    # Run an approved DAG; completed checkpoints are skipped on resume.
    #{self}.execute(
      dag: 'optional - Hash DAG; required unless path or run_id is set',
      path: 'optional - YAML DAG path',
      run_id: 'optional - existing run to resume; generated when omitted',
      root: 'optional - runs directory; defaults to ~/.pwn/runs',
      halt_after: 'optional - step id used by tests to stop after a checkpoint',
      approved: 'optional - true allows shell steps on an unattended mission',
      unattended: 'optional - true refuses inferred shell steps unless approved',
      operator: 'optional - true rejects inferred shell steps even when approved',
      mission_id: 'optional - mission that records a hand-written shell exception'
    )

    # Resume a checkpointed run, skipping completed steps.
    #{self}.resume(
      run_id: 'required - run id previously returned by execute',
      root: 'optional - runs directory; defaults to ~/.pwn/runs'
    )

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

.plan(opts = {}) ⇒ Object

Emit a YAML-serializable DAG with zero tool side effects.



21
22
23
24
25
26
27
28
29
30
31
32
 
# File 'lib/pwn/ai/agent/task_dag.rb', line 21

public_class_method def self.plan(opts = {})
  request = opts[:request].to_s
  raise 'ERROR: request is required' if request.strip.empty?

  steps = Array(opts[:steps])
  steps = infer_steps(request: request) if steps.empty?
  {
    version: 1,
    request: request,
    steps: steps.map { |step| normalize_step(step: step) }
  }
end

.required_bins ⇒ Object 



16
17
18
 
# File 'lib/pwn/ai/agent/task_dag.rb', line 16

public_class_method def self.required_bins
  []
end

.resume(opts = {}) ⇒ Object

Resume a checkpointed run, skipping completed steps.



66
67
68
 
# File 'lib/pwn/ai/agent/task_dag.rb', line 66

public_class_method def self.resume(opts = {})
  execute(opts)
end