iop/agent-ops/skills/project/openai-usage-token-issue/scripts/remote_edge_transaction.rb

# frozen_string_literal: true

require "digest"
require "fileutils"
require "json"
require "net/http"
require "open3"
require "securerandom"
require "time"
require "tmpdir"
require "yaml"

NOHUP_BIN = "/usr/bin/nohup"
PROFILES = {
  "dev" => {
    "environment" => "dev",
    "config_path" => "/Users/toki/agent-work/iop-dev/build/dev-runtime/edge.yaml",
    "config_arg" => "build/dev-runtime/edge.yaml",
    "edge_bin" => "/Users/toki/agent-work/iop-dev/build/dev-runtime/bin/edge",
    "runtime_root" => "/Users/toki/agent-work/iop-dev",
    "admin_addr" => "127.0.0.1:19093",
    "health_url" => "http://127.0.0.1:18083/healthz",
    "listener_port" => 18_083,
    "log_path" => "/Users/toki/agent-work/iop-dev/build/dev-runtime/logs/iop-edge-token-issue.log",
    "lock_path" => "/Users/toki/agent-work/iop-dev/build/dev-runtime/.token-issue.lock",
    "api_base_url_runner" => "http://127.0.0.1:18083/v1",
    "metrics_url_runner" => "http://127.0.0.1:19101/metrics",
    "smoke_model" => "laguna-s:2.1"
  },
  "dev-corp" => {
    "environment" => "dev-corp",
    "config_path" => "/Users/toki/agent-work/iop-dev-corp/build/dev-corp-runtime/edge.yaml",
    "config_arg" => "/Users/toki/agent-work/iop-dev-corp/build/dev-corp-runtime/edge.yaml",
    "edge_bin" => "/Users/toki/agent-work/iop-dev-corp/build/dev-corp-runtime/bin/iop-edge",
    "runtime_root" => "/Users/toki/agent-work/iop-dev-corp",
    "admin_addr" => "127.0.0.1:19094",
    "health_url" => "http://127.0.0.1:18086/healthz",
    "listener_port" => 18_086,
    "log_path" => "/Users/toki/agent-work/iop-dev-corp/build/dev-corp-runtime/logs/iop-edge-token-issue.log",
    "lock_path" => "/Users/toki/agent-work/iop-dev-corp/build/dev-corp-runtime/.token-issue.lock",
    "api_base_url_runner" => "http://127.0.0.1:18086/v1",
    "metrics_url_runner" => "http://127.0.0.1:19105/metrics",
    "smoke_model" => "ornith:35b"
  }
}.freeze

def profile
  $profile || fail_safe("environment_missing")
end

def config_path
  profile.fetch("config_path")
end

def config_arg
  profile.fetch("config_arg")
end

def edge_bin
  profile.fetch("edge_bin")
end

def runtime_root
  profile.fetch("runtime_root")
end

def admin_addr
  profile.fetch("admin_addr")
end

def health_uri
  URI(profile.fetch("health_url"))
end

def listener_port
  Integer(profile.fetch("listener_port"))
end

def log_path
  profile.fetch("log_path")
end

def lock_path
  profile.fetch("lock_path")
end

def expected_command
  "#{edge_bin} --config #{config_arg} serve"
end

class SafeFailure < StandardError
  attr_reader :code

  def initialize(code)
    @code = code
    super(code)
  end
end

def fail_safe(code)
  raise SafeFailure, code
end

def with_transaction_lock
  File.open(lock_path, File::RDWR | File::CREAT, 0o600) do |lock|
    File.chmod(0o600, lock_path)
    fail_safe("transaction_busy") unless lock.flock(File::LOCK_EX | File::LOCK_NB)
    begin
      yield
    ensure
      lock.flock(File::LOCK_UN)
    end
  end
end

def read_payload
  payload = JSON.parse($stdin.read)
  fail_safe("invalid_action") unless %w[inspect candidate-check apply rollback api-smoke metrics selftest].include?(payload["action"])
  environment = payload["environment"]
  fail_safe("environment_invalid") unless environment.is_a?(String) && PROFILES.key?(environment)
  payload
rescue JSON::ParserError
  fail_safe("invalid_json")
end

def load_config(path = config_path)
  value = YAML.safe_load(File.read(path), aliases: true)
  fail_safe("invalid_config_root") unless value.is_a?(Hash)
  value
rescue Errno::ENOENT
  fail_safe("config_missing")
rescue Psych::Exception
  fail_safe("config_yaml_invalid")
end

def mappings(config)
  items = config.dig("openai", "principal_tokens") || []
  fail_safe("principal_tokens_invalid") unless items.is_a?(Array) && items.all? { |item| item.is_a?(Hash) }
  items
end

def validate_mappings!(items)
  refs = items.map { |item| item["token_ref"].to_s }
  hashes = items.map { |item| item["token_hash_sha256"].to_s.downcase }
  fail_safe("mapping_token_ref_invalid") if refs.any?(&:empty?) || refs.uniq.length != refs.length
  fail_safe("mapping_hash_invalid") if hashes.any? { |value| !value.match?(/\A[0-9a-f]{64}\z/) } || hashes.uniq.length != hashes.length
  fail_safe("mapping_principal_invalid") if items.any? { |item| item["principal_ref"].to_s.empty? }
end

def json_scalar(value)
  JSON.generate(value.to_s)
end

def insert_mapping(source, entry)
  lines = source.lines
  openai_index = lines.index { |line| line.match?(/\Aopenai:\s*(?:#.*)?\z/) }
  fail_safe("openai_section_missing") unless openai_index

  openai_end = ((openai_index + 1)...lines.length).find do |index|
    lines[index].match?(/\A\S/) && !lines[index].lstrip.start_with?("#")
  end || lines.length
  token_index = ((openai_index + 1)...openai_end).find do |index|
    lines[index].match?(/\A  principal_tokens:\s*(?:\[\])?\s*(?:#.*)?\z/)
  end
  fail_safe("principal_tokens_section_missing") unless token_index

  if lines[token_index].match?(/principal_tokens:\s*\[\]/)
    lines[token_index] = "  principal_tokens:\n"
    token_end = token_index + 1
    sequence_indent = 4
  else
    token_end = ((token_index + 1)...openai_end).find do |index|
      lines[index].match?(/\A  [A-Za-z0-9_][A-Za-z0-9_-]*:/)
    end || openai_end
    first_item = lines[(token_index + 1)...token_end].find { |line| line.match?(/\A\s*-\s+token_ref:/) }
    fail_safe("principal_tokens_style_invalid") unless first_item
    sequence_indent = first_item[/\A\s*/].length
  end
  sequence_prefix = " " * sequence_indent
  field_prefix = " " * (sequence_indent + 2)
  fragment = [
    "#{sequence_prefix}- token_ref: #{json_scalar(entry.fetch("token_ref"))}\n",
    "#{field_prefix}token_hash_sha256: #{json_scalar(entry.fetch("token_hash_sha256"))}\n",
    "#{field_prefix}principal_ref: #{json_scalar(entry.fetch("principal_ref"))}\n",
    "#{field_prefix}principal_alias: #{json_scalar(entry.fetch("principal_alias"))}\n"
  ]
  lines.insert(token_end, *fragment)
  lines.join
end

def run_command(*argv)
  stdout, _stderr, status = Open3.capture3(*argv)
  fail_safe("command_failed") unless status.success?
  stdout
end

def loopback_uri(profile_key, path = nil)
  base = profile.fetch(profile_key)
  uri = URI(path ? "#{base}#{path}" : base)
  fail_safe("loopback_url_invalid") unless uri.scheme == "http" && %w[127.0.0.1 ::1 localhost].include?(uri.host)
  fail_safe("loopback_url_invalid") if uri.user || uri.password || uri.fragment
  uri
end

def loopback_request(uri, request, read_timeout:)
  http = Net::HTTP.new(uri.host, uri.port, nil)
  http.open_timeout = 3
  http.read_timeout = read_timeout
  response = http.start { |client| client.request(request) }
  fail_safe("loopback_http_failed") unless response.is_a?(Net::HTTPSuccess)
  response.body.to_s
rescue Net::OpenTimeout, Net::ReadTimeout, SocketError, SystemCallError
  fail_safe("loopback_network_failed")
end

def api_json(path, raw_token, payload = nil, read_timeout: 30)
  uri = loopback_uri("api_base_url_runner", path)
  request = payload ? Net::HTTP::Post.new(uri) : Net::HTTP::Get.new(uri)
  request["Authorization"] = "Bearer #{raw_token}"
  request["Content-Type"] = "application/json"
  request.body = JSON.generate(payload) if payload
  body = loopback_request(uri, request, read_timeout: read_timeout)
  parsed = JSON.parse(body)
  fail_safe("openai_response_invalid") unless parsed.is_a?(Hash)
  parsed
rescue JSON::ParserError
  fail_safe("openai_response_invalid")
end

def api_smoke(payload)
  raw_token = payload["raw_token"]
  fail_safe("raw_token_invalid") unless raw_token.is_a?(String) && raw_token.start_with?("iop_") && !raw_token.match?(/[\r\n]/)
  models = api_json("/models", raw_token, nil, read_timeout: 15)
  fail_safe("openai_models_invalid") unless models["data"].is_a?(Array) && !models["data"].empty?
  response = api_json(
    "/chat/completions",
    raw_token,
    {
      "model" => profile.fetch("smoke_model"),
      "messages" => [{ "role" => "user", "content" => "Reply with the single word OK." }],
      "max_tokens" => 2048,
      "temperature" => 0
    },
    read_timeout: 120
  )
  choice = response["choices"].is_a?(Array) ? response["choices"].first : nil
  message = choice.is_a?(Hash) ? choice["message"] : nil
  content = message.is_a?(Hash) ? message["content"] : nil
  fail_safe("openai_chat_invalid") unless content.is_a?(String) && !content.strip.empty? && choice["finish_reason"]
  { "status" => "passed" }
end

def metrics_observed(payload)
  token_ref = payload["token_ref"]
  fail_safe("token_ref_invalid") unless token_ref.is_a?(String) && token_ref.match?(/\A[a-z0-9][a-z0-9._:-]{2,79}\z/)
  uri = loopback_uri("metrics_url_runner")
  body = loopback_request(uri, Net::HTTP::Get.new(uri), read_timeout: 8)
  { "status" => "ok", "observed" => body.match?(/token_ref="#{Regexp.escape(token_ref)}"/) }
end

def parse_refresh_status(stdout)
  parsed = begin
    value = JSON.parse(stdout)
    value if value.is_a?(Hash) && value.key?("status")
  rescue JSON::ParserError
    nil
  end
  unless parsed
    stdout.lines.reverse_each do |line|
      begin
        value = JSON.parse(line)
        if value.is_a?(Hash) && value.key?("status")
          parsed = value
          break
        end
      rescue JSON::ParserError
        next
      end
    end
  end
  fail_safe("refresh_response_invalid") unless parsed
  parsed.fetch("status").to_s
end

def refresh_status(candidate_path)
  stdout = run_command(
    edge_bin,
    "--config", config_path,
    "config", "refresh",
    "--addr", admin_addr,
    "--config-path", candidate_path,
    "--mode", "dry-run"
  )
  parse_refresh_status(stdout)
end

def listener_pid
  stdout, _stderr, status = Open3.capture3("lsof", "-tiTCP:#{listener_port}", "-sTCP:LISTEN")
  return nil unless status.success?
  values = stdout.lines.map(&:strip).reject(&:empty?).uniq
  fail_safe("edge_listener_ambiguous") unless values.length == 1
  value = values.first
  fail_safe("edge_listener_pid_invalid") unless value.match?(/\A\d+\z/)
  value.to_i
end

def expected_listener_pid
  pid = listener_pid
  return nil unless pid
  stdout, _stderr, status = Open3.capture3("ps", "-p", pid.to_s, "-o", "command=")
  fail_safe("edge_listener_identity_unavailable") unless status.success?
  fail_safe("edge_listener_unexpected") unless stdout.strip == expected_command
  pid
end

def healthy?
  uri = health_uri
  response = Net::HTTP.start(uri.host, uri.port, open_timeout: 2, read_timeout: 2) do |http|
    http.get(uri.request_uri)
  end
  response.is_a?(Net::HTTPSuccess)
rescue StandardError
  false
end

def wait_until(timeout_seconds)
  deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout_seconds
  loop do
    return true if yield
    return false if Process.clock_gettime(Process::CLOCK_MONOTONIC) >= deadline
    sleep 0.25
  end
end

def stop_edge
  pid = expected_listener_pid
  return unless pid

  Process.kill("TERM", pid)
  return if wait_until(10) { expected_listener_pid != pid }

  Process.kill("KILL", pid)
  fail_safe("edge_stop_failed") unless wait_until(5) { expected_listener_pid != pid }
rescue Errno::ESRCH
  nil
rescue Errno::EPERM
  fail_safe("edge_stop_denied")
end

def start_edge
  FileUtils.mkdir_p(File.dirname(log_path))
  log = File.open(log_path, "a", 0o600)
  File.chmod(0o600, log_path)
  begin
    pid = Process.spawn(
      NOHUP_BIN, edge_bin, "--config", config_arg, "serve",
      chdir: runtime_root,
      in: File::NULL,
      out: log,
      err: log,
      pgroup: true
    )
  ensure
    log.close
  end
  Process.detach(pid)
  fail_safe("edge_readiness_failed") unless wait_until(30) { expected_listener_pid == pid && healthy? }
end

def restart_edge
  stop_edge
  start_edge
end

def cutover_files(config_path, candidate_path, backup_path)
  File.rename(config_path, backup_path)
  begin
    File.rename(candidate_path, config_path)
  rescue StandardError
    File.rename(backup_path, config_path) if File.file?(backup_path) && !File.exist?(config_path)
    raise
  end
end

def rollback_to(backup_path)
  expanded = File.expand_path(backup_path.to_s)
  allowed_prefix = File.dirname(config_path) + File::SEPARATOR
  fail_safe("backup_path_invalid") unless expanded.start_with?(allowed_prefix)
  fail_safe("backup_path_invalid") unless File.basename(expanded).match?(/\Aedge\.yaml\.before-token-\d{8}T\d{6}Z-[0-9a-f]{6}\.yaml\z/)
  fail_safe("backup_missing") unless File.file?(expanded)
  run_command(edge_bin, "--config", expanded, "config", "check")

  failed_path = "#{config_path}.failed-token-#{Time.now.utc.strftime("%Y%m%dT%H%M%SZ")}-#{SecureRandom.hex(3)}.yaml"
  stop_edge
  File.rename(config_path, failed_path) if File.exist?(config_path)
  File.rename(expanded, config_path)
  start_edge
  fail_safe("rollback_config_invalid") unless refresh_status(config_path) == "applied"
  { "status" => "rolled_back" }
end

def inspect_state(payload)
  fail_safe("runtime_assets_invalid") unless File.executable?(edge_bin) && File.executable?(NOHUP_BIN)
  fail_safe("runtime_directory_not_writable") unless File.writable?(File.dirname(config_path))
  config = load_config
  items = mappings(config)
  validate_mappings!(items)
  probe_entry = {
    "token_ref" => "iop-#{profile.fetch("environment")}-preflight-#{SecureRandom.hex(6)}",
    "token_hash_sha256" => Digest::SHA256.hexdigest(SecureRandom.random_bytes(32)),
    "principal_ref" => "preflight-shape.invalid",
    "principal_alias" => "preflight-shape"
  }
  probe_config = YAML.safe_load(insert_mapping(File.read(config_path), probe_entry), aliases: true)
  probe_items = mappings(probe_config)
  validate_mappings!(probe_items)
  fail_safe("candidate_shape_invalid") unless probe_items.length == items.length + 1
  current_refresh = payload["refresh_probe"] ? refresh_status(config_path) : nil
  fail_safe("current_refresh_invalid") if current_refresh && current_refresh != "applied"
  {
    "status" => "ok",
    "mapping_count" => items.length,
    "mappings" => items.map do |item|
      {
        "token_ref" => item["token_ref"].to_s,
        "token_hash_sha256" => item["token_hash_sha256"].to_s.downcase,
        "principal_ref" => item["principal_ref"].to_s,
        "principal_alias" => item["principal_alias"].to_s
      }
    end,
    "candidate_shape" => true,
    "refresh_status" => current_refresh,
    "healthy" => healthy?,
    "listener" => !expected_listener_pid.nil?
  }
end

def candidate_check
  active_items = mappings(load_config)
  probe_entry = {
    "token_ref" => "iop-#{profile.fetch("environment")}-candidate-check-#{SecureRandom.hex(6)}",
    "token_hash_sha256" => Digest::SHA256.hexdigest(SecureRandom.random_bytes(32)),
    "principal_ref" => "candidate-check.invalid",
    "principal_alias" => "candidate-check"
  }
  candidate_path = File.join(File.dirname(config_path), ".edge-token-candidate-check-#{SecureRandom.hex(6)}.yaml")
  begin
    File.open(candidate_path, File::WRONLY | File::CREAT | File::EXCL, 0o600) do |file|
      file.write(insert_mapping(File.read(config_path), probe_entry))
      file.flush
      file.fsync
    end
    candidate_items = mappings(load_config(candidate_path))
    validate_mappings!(candidate_items)
    fail_safe("candidate_count_invalid") unless candidate_items.length == active_items.length + 1
    run_command(edge_bin, "--config", candidate_path, "config", "check")
    fail_safe("restart_required_not_reported") unless refresh_status(candidate_path) == "restart_required"
    { "status" => "validated" }
  ensure
    File.delete(candidate_path) if File.exist?(candidate_path)
  end
end

def apply(payload)
  entry = payload.fetch("entry")
  fail_safe("entry_invalid") unless entry.is_a?(Hash)
  fail_safe("hash_invalid") unless entry["token_hash_sha256"].to_s.match?(/\A[0-9a-f]{64}\z/)
  %w[token_ref principal_ref principal_alias].each do |key|
    value = entry[key].to_s
    fail_safe("entry_invalid") if value.empty? || value.include?("\n") || value.include?("\r")
  end

  active = load_config
  active_items = mappings(active)
  validate_mappings!(active_items)
  fail_safe("principal_conflict") if active_items.any? { |item| item["principal_ref"].to_s == entry["principal_ref"].to_s }
  fail_safe("token_ref_conflict") if active_items.any? { |item| item["token_ref"].to_s == entry["token_ref"].to_s }
  fail_safe("token_hash_conflict") if active_items.any? { |item| item["token_hash_sha256"].to_s.casecmp?(entry["token_hash_sha256"].to_s) }

  candidate_path = nil
  candidate_items = nil
  backup_path = nil
  original_moved = false
  begin
    source = File.read(config_path)
    candidate_path = File.join(
      File.dirname(config_path),
      ".edge-token-issue-#{Time.now.utc.strftime("%Y%m%dT%H%M%SZ")}-#{SecureRandom.hex(3)}.yaml"
    )
    File.open(candidate_path, File::WRONLY | File::CREAT | File::EXCL, 0o600) do |file|
      file.write(insert_mapping(source, entry))
      file.flush
      file.fsync
    end

    candidate = load_config(candidate_path)
    candidate_items = mappings(candidate)
    validate_mappings!(candidate_items)
    fail_safe("candidate_count_invalid") unless candidate_items.length == active_items.length + 1
    run_command(edge_bin, "--config", candidate_path, "config", "check")
    fail_safe("restart_required_not_reported") unless refresh_status(candidate_path) == "restart_required"

    backup_path = "#{config_path}.before-token-#{Time.now.utc.strftime("%Y%m%dT%H%M%SZ")}-#{SecureRandom.hex(3)}.yaml"
    cutover_files(config_path, candidate_path, backup_path)
    original_moved = true
    candidate_path = nil
    restart_edge
    current = mappings(load_config)
    validate_mappings!(current)
    expected = current.count do |item|
      item["principal_ref"].to_s == entry["principal_ref"].to_s &&
        item["token_ref"].to_s == entry["token_ref"].to_s &&
        item["token_hash_sha256"].to_s.casecmp?(entry["token_hash_sha256"].to_s)
    end
    fail_safe("active_mapping_invalid") unless expected == 1
    fail_safe("post_restart_refresh_invalid") unless refresh_status(config_path) == "applied"
  rescue StandardError
    if original_moved && backup_path && File.file?(backup_path)
      begin
        rollback_to(backup_path)
      rescue StandardError
        fail_safe("automatic_rollback_failed")
      end
    end
    raise
  ensure
    File.delete(candidate_path) if candidate_path && File.exist?(candidate_path)
  end

  {
    "status" => "activated",
    "backup_path" => backup_path,
    "mapping_count" => candidate_items.length
  }
end

def selftest
  PROFILES.each_value do |item|
    fail_safe("selftest_profile_path_failed") unless item.fetch("config_path").start_with?(item.fetch("runtime_root") + File::SEPARATOR)
    fail_safe("selftest_profile_command_failed") unless item.fetch("edge_bin").start_with?(item.fetch("runtime_root") + File::SEPARATOR)
    uri = URI(item.fetch("api_base_url_runner"))
    fail_safe("selftest_profile_url_failed") unless uri.scheme == "http" && uri.host == "127.0.0.1"
  end
  sample = <<~YAML
    version: 1
    openai:
      enabled: true
      principal_tokens: []
      timeout_sec: 30
  YAML
  entry = {
    "token_ref" => "iop-dev-corp-sample",
    "token_hash_sha256" => "a" * 64,
    "principal_ref" => "sample.invalid",
    "principal_alias" => "sample"
  }
  inserted = YAML.safe_load(insert_mapping(sample, entry), aliases: true)
  items = mappings(inserted)
  validate_mappings!(items)
  fail_safe("selftest_insert_failed") unless items.length == 1
  indentless = <<~YAML
    version: 1
    openai:
      enabled: true
      principal_tokens:
      - token_ref: "existing"
        token_hash_sha256: "#{"b" * 64}"
        principal_ref: "existing.invalid"
        principal_alias: "existing"
    a2a:
      enabled: false
  YAML
  indentless_inserted = insert_mapping(indentless, entry)
  fail_safe("selftest_indentless_style_failed") unless indentless_inserted.include?("\n  - token_ref: \"iop-dev-corp-sample\"\n")
  indentless_items = mappings(YAML.safe_load(indentless_inserted, aliases: true))
  validate_mappings!(indentless_items)
  fail_safe("selftest_indentless_insert_failed") unless indentless_items.length == 2
  applied = JSON.pretty_generate({ "status" => "applied", "restart_required_paths" => ["openai.principal_tokens"] })
  fail_safe("selftest_status_failed") unless parse_refresh_status(applied) == "applied"
  Dir.mktmpdir("iop-token-cutover-") do |directory|
    config_path = File.join(directory, "edge.yaml")
    missing_candidate = File.join(directory, "missing.yaml")
    backup_path = File.join(directory, "edge.yaml.backup")
    File.write(config_path, "original")
    begin
      cutover_files(config_path, missing_candidate, backup_path)
      fail_safe("selftest_cutover_failed")
    rescue Errno::ENOENT
      fail_safe("selftest_cutover_restore_failed") unless File.read(config_path) == "original"
      fail_safe("selftest_cutover_backup_leaked") if File.exist?(backup_path)
    end
  end
  { "status" => "ok" }
end

begin
  payload = read_payload
  $profile = PROFILES.fetch(payload.fetch("environment"))
  result = case payload.fetch("action")
           when "inspect" then inspect_state(payload)
           when "candidate-check" then with_transaction_lock { candidate_check }
           when "apply" then with_transaction_lock { apply(payload) }
           when "rollback" then with_transaction_lock { rollback_to(payload.fetch("backup_path")) }
           when "api-smoke" then api_smoke(payload)
           when "metrics" then metrics_observed(payload)
           when "selftest" then selftest
           end
  puts JSON.generate(result)
rescue SafeFailure => error
  puts JSON.generate({ "status" => "blocked", "code" => error.code })
  exit 2
rescue KeyError
  puts JSON.generate({ "status" => "blocked", "code" => "missing_input" })
  exit 2
rescue StandardError
  puts JSON.generate({ "status" => "blocked", "code" => "unexpected_remote_failure" })
  exit 2
end