trmnl_preview 0.7.1 → 0.8.0

data/lib/trmnlp/screen_generator.rb

@@ -1,238 +1,32 @@
-require 'mini_magick'
-require 'selenium-webdriver'
-require 'base64'
-require 'thread'
-require 'tempfile'
-require 'fileutils'
-require 'uri'
+# frozen_string_literal: true
+
+require_relative 'screenshot'
+require_relative 'image_quantizer'
 
 module TRMNLP
   class ScreenGenerator
-    # Browser pool management for efficient resource usage
-    class BrowserPool
-      def initialize(max_size: 2)
-        @drivers = []
-        @available = Queue.new
-        @mutex = Mutex.new
-        @max_size = max_size
-        @shutdown = false
-
-        at_exit { shutdown }
-      end
-
-      def with_driver
-        driver = nil
-
-        begin
-          driver = checkout_driver
-          yield driver
-        ensure
-          checkin_driver(driver) if driver
-        end
-      end
-
-      def shutdown
-        @mutex.synchronize do
-          return if @shutdown
-          @shutdown = true
-
-          @drivers.each do |driver|
-            driver.quit rescue nil
-          end
-
-          @drivers.clear
-        end
-      end
-
-      private
-
-      def checkout_driver
-        driver = @available.pop(true) rescue nil
-
-        if driver.nil?
-          @mutex.synchronize do
-            if @drivers.size < @max_size
-              driver = create_driver
-              @drivers << driver
-            end
-          end
-        end
-
-        driver ||= @available.pop
-
-        begin
-          # Ping the driver
-          driver.title
-          driver
-        rescue
-          @mutex.synchronize do
-            @drivers.delete(driver)
-            driver = create_driver
-            @drivers << driver
-          end
-          driver
-        end
-      end
-
-      def checkin_driver(driver)
-        return if @shutdown
-        @available.push(driver)
-      end
-
-      def create_driver
-        options = Selenium::WebDriver::Firefox::Options.new
-        options.add_argument('--headless')
-        options.add_argument('--disable-web-security')
-
-        driver = Selenium::WebDriver.for(:firefox, options: options)
-        # Set a default window size that will be consistent
-        driver.manage.window.maximize
-        driver
-      end
-    end
-
-    @@browser_pool = BrowserPool.new
-
     def initialize(html, opts = {})
-      self.input = html
-      self.image = !!opts[:image]
-
-      # Accept optional rendering parameters (width/height/color depth/dark mode)
+      @input = html
+      @screenshot = opts[:screenshot]
       @requested_width = opts[:width]
       @requested_height = opts[:height]
       @requested_color_depth = opts[:color_depth]
     end
 
-    attr_accessor :input, :output, :image
-
     def process
-      convert_to_image
-      image ? mono_image(output) : mono(output)
+      output = @screenshot.call(html: @input, width:, height:)
+      ImageQuantizer.new(depth: color_depth).call(output.path)
       output
     end
 
     private
 
-    def convert_to_image
-      retry_count = 0
-
-      begin
-        @@browser_pool.with_driver do |driver|
-          # determine dimensions of toolbars, etc
-          borders = driver.execute_script(<<~JS)
-            return {
-              width: window.outerWidth - window.innerWidth,
-              height: window.outerHeight - window.innerHeight
-            }
-          JS
-
-          window_width = width + borders['width']
-          window_height = height + borders['height']
-          driver.manage.window.size = Selenium::WebDriver::Dimension.new(window_width, window_height)
-          
-          sleep(0.1)
-
-          prepare_page(driver)
-
-          self.output = Tempfile.new(['screenshot', '.png'])
-          driver.save_screenshot(output.path)
-          output.close
-        end
-      rescue Selenium::WebDriver::Error::TimeoutError,
-            Selenium::WebDriver::Error::WebDriverError => e
-        retry_count += 1
-        retry if retry_count <= 1
-        raise
-      end
-    end
-
-    def prepare_page(driver)
-      driver.navigate.to('about:blank')
-
-      driver.execute_script(<<~JS, input)
-        document.open();
-        document.write(arguments[0]);
-        document.close();
-      JS
-
-      Selenium::WebDriver::Wait.new(timeout: 5).until do
-        driver.execute_script('return document.readyState') == 'complete'
-      end
-
-      # Wait for fonts (prevents layout shifts)
-      driver.execute_script('return document.fonts && document.fonts.ready')
-
-      driver.execute_script(<<~JS)
-        document.documentElement.style.overflow = 'hidden';
-        document.body.style.overflow = 'hidden';
-      JS
-    end
-      
-    def convert_with_mini_magick(img, depth)
-      tmp = Tempfile.new(['mono', '.png'])
-      tmp.close
-
-      levels = 2**depth
-
-      MiniMagick::Tool::Convert.new do |m|
-        m << img.path
-        m.colorspace 'Gray'
-        m.dither 'FloydSteinberg'
-
-        yield(m, depth, levels)
-
-        m.depth depth
-        m.define "png:bit-depth=#{depth}"
-        m.strip
-        m << tmp.path
-      end
-
-      FileUtils.mv(tmp.path, img.path, force: true)
-    end
-
-    def mono(img)
-      depth = [[color_depth.to_i, 1].max, 8].min
-
-      convert_with_mini_magick(img, depth) do |m, d, levels|
-        m.posterize levels
-        m.colors levels
-        m.type 'Bilevel' if d == 1
-      end
-    end
-
-    def mono_image(img)
-      depth = [[color_depth.to_i, 1].max, 8].min
-
-      convert_with_mini_magick(img, depth) do |m, d, levels|
-        if d == 1
-          # For true 1-bit, use a halftone/remap and bilevel output
-          m.remap 'pattern:gray50'
-          m.posterize 2
-          m.colors 2
-          m.type 'Bilevel'
-        else
-          m.posterize levels
-          m.colors levels
-        end
-      end
-    end
-
-
-    def width
-      @requested_width || 800
-    end
-
-    def height
-      @requested_height || 480
-    end
+    def width = @requested_width || 800
+    def height = @requested_height || 480
 
     def color_depth
       return @requested_color_depth if @requested_color_depth
-
-      # Try to infer color depth from the rendered HTML's screen classes
-      if input && input.match(/screen--(\d+)bit/)
-        return $1.to_i
-      end
+      return ::Regexp.last_match(1).to_i if @input&.match(/screen--(\d+)bit/)
 
       1
     end

data/lib/trmnlp/screenshot.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+require 'selenium-webdriver'
+require 'tempfile'
+
+module TRMNLP
+  class Screenshot
+    def initialize(pool:)
+      @pool = pool
+    end
+
+    def call(html:, width:, height:)
+      attempts = 0
+
+      begin
+        @pool.with_driver { |driver| render(driver, html, width, height) }
+      rescue Selenium::WebDriver::Error::TimeoutError,
+             Selenium::WebDriver::Error::WebDriverError
+        attempts += 1
+        retry if attempts <= 1
+        raise
+      end
+    end
+
+    private
+
+    def render(driver, html, width, height)
+      resize(driver, width, height)
+      load_page(driver, html)
+      capture(driver)
+    end
+
+    def resize(driver, width, height)
+      apply_window_size(driver, width, height)
+      wait_for_viewport(driver, width, height)
+    end
+
+    def apply_window_size(driver, width, height)
+      borders = driver.execute_script(<<~JS)
+        return {
+          width: window.outerWidth - window.innerWidth,
+          height: window.outerHeight - window.innerHeight
+        }
+      JS
+
+      dim = Selenium::WebDriver::Dimension.new(width + borders['width'], height + borders['height'])
+      driver.manage.window.size = dim
+    end
+
+    # NOTE: A cold Firefox — e.g. the first render after the container boots —
+    # applies a window resize lazily. The old fixed sleep raced that reflow and
+    # clipped the first screenshot short (800x433 instead of 800x480). Poll the
+    # real viewport instead, re-applying the size until it lands; a resize that
+    # never settles surfaces as a TimeoutError that #call already retries.
+    def wait_for_viewport(driver, width, height)
+      Selenium::WebDriver::Wait.new(timeout: 5, interval: 0.1).until do
+        next true if viewport(driver) == [width, height]
+
+        apply_window_size(driver, width, height)
+        false
+      end
+    end
+
+    def viewport(driver)
+      driver.execute_script('return [window.innerWidth, window.innerHeight]')
+    end
+
+    def load_page(driver, html)
+      driver.navigate.to('about:blank')
+
+      driver.execute_script(<<~JS, html)
+        document.open();
+        document.write(arguments[0]);
+        document.close();
+      JS
+
+      Selenium::WebDriver::Wait.new(timeout: 5).until do
+        driver.execute_script('return document.readyState') == 'complete'
+      end
+
+      driver.execute_script('return document.fonts && document.fonts.ready')
+
+      driver.execute_script(<<~JS)
+        document.documentElement.style.overflow = 'hidden';
+        document.body.style.overflow = 'hidden';
+      JS
+    end
+
+    def capture(driver)
+      file = Tempfile.new(['screenshot', '.png'])
+      driver.save_screenshot(file.path)
+      file.close
+      file
+    end
+  end
+end

data/lib/trmnlp/transform_backend/http.rb

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+require 'faraday'
+require 'json'
+
+require_relative '../transform_client'
+require_relative 'wrapper'
+
+module TRMNLP
+  module TransformBackend
+    # Remote-daemon transform execution. Speaks the production daemon's
+    # wire format over HTTP so trmnlp can target a real remote transform
+    # daemon (or any compatible server) instead of running transforms
+    # locally. Selected by TransformClient.from_config when
+    # serverless_daemon_url is set in the project's .trmnlp.yml.
+    #
+    # The daemon expects code that already includes its own harness
+    # (reading stdin, dispatching to run/transform/result, writing the
+    # canonical JSON result to FD 3). The shared Wrapper module emits
+    # the same harness Subprocess uses, parameterized on a FD-3 sink
+    # so a transform behaves identically whether previewed against the
+    # daemon or run locally.
+    class Http
+      SUPPORTED_LANGUAGES = %w[python ruby php node].freeze
+      DEFAULT_TIMEOUT = 30
+      HTTP_TIMEOUT = 60
+
+      def initialize(url:, api_key: nil, http_timeout: HTTP_TIMEOUT)
+        @url = url
+        @api_key = api_key
+        @http_timeout = http_timeout
+      end
+
+      def execute(code:, language:, stdin: '', timeout_seconds: DEFAULT_TIMEOUT)
+        lang = language.to_s
+        return failure("unsupported serverless_language: #{lang}") unless SUPPORTED_LANGUAGES.include?(lang)
+
+        post(code: Wrapper.for(lang, code, sink_for(lang)), stdin:, timeout: timeout_seconds, language: lang)
+      end
+
+      private
+
+      def post(code:, stdin:, timeout:, language:)
+        response = connection.post('/execute') do |req|
+          req.body = JSON.generate(code:, stdin:, timeout:, language:)
+        end
+
+        return failure("daemon HTTP #{response.status}: #{response.body}") unless response.status == 200
+
+        parse(response.body)
+      rescue Faraday::ConnectionFailed, Faraday::TimeoutError => e
+        failure("transform daemon unreachable at #{@url}: #{e.message}")
+      rescue JSON::ParserError => e
+        failure("daemon returned non-JSON response: #{e.message}")
+      end
+
+      def parse(body)
+        parsed = JSON.parse(body)
+        TransformClient::Result.new(
+          stdout: parsed['stdout'] || '',
+          stderr: parsed['stderr'] || '',
+          # Fall back to stdout for daemons that haven't been upgraded to
+          # the separate `output` channel yet.
+          output: (parsed['output'].to_s.empty? ? parsed['stdout'] : parsed['output']).to_s,
+          exit_code: parsed['exit_code'] || 0,
+          duration_ms: parsed['duration_ms'] || 0,
+          error: parsed['error']
+        )
+      end
+
+      # NOTE: the connection is memoized and never explicitly closed. That is
+      # safe only because Faraday's default adapter opens a fresh socket per
+      # request — swapping in a persistent adapter here would leak sockets.
+      def connection
+        @connection ||= Faraday.new(url: @url) do |f|
+          f.headers['Content-Type'] = 'application/json'
+          f.headers['Authorization'] = "Bearer #{@api_key}" if @api_key
+          f.options.timeout = @http_timeout
+          f.options.open_timeout = 5
+        end
+      end
+
+      def failure(message)
+        TransformClient::Result.new(stdout: '', stderr: '', output: '', exit_code: -1, duration_ms: 0, error: message)
+      end
+
+      # Language-specific FD-3 sink snippets — what the daemon's
+      # harness reads from to capture canonical JSON output.
+      def sink_for(language)
+        case language
+        when 'python'
+          "os.write(3, json.dumps(output).encode('utf-8'))"
+        when 'ruby'
+          'IO.new(3).write(JSON.generate(output))'
+        when 'node'
+          <<~JS.chomp
+            Promise.resolve(output).then(o => {
+              require('fs').writeSync(3, JSON.stringify(o));
+            });
+          JS
+        when 'php'
+          "$fd = fopen('php://fd/3', 'w');\n          fwrite($fd, json_encode($output));\n          fclose($fd);"
+        end
+      end
+    end
+  end
+end

data/lib/trmnlp/transform_backend/subprocess.rb

@@ -0,0 +1,130 @@
+# frozen_string_literal: true
+
+require 'open3'
+require 'tmpdir'
+
+require_relative '../transform_client'
+require_relative 'wrapper'
+
+module TRMNLP
+  module TransformBackend
+    # Local subprocess execution of user transform code. Mirrors the
+    # remote-daemon wrapper contract so a transform behaves the same
+    # locally as it does in production. Output flows back via
+    # a tempfile per-execution instead of FD 3, but the
+    # run/result/input dispatch logic is preserved verbatim via the
+    # shared Wrapper module.
+    class Subprocess
+      DEFAULT_TIMEOUT = 30
+      # Seconds a TERM'd process is given to exit before escalating to KILL.
+      GRACE_PERIOD = 0.1
+
+      INTERPRETERS = {
+        'python' => { cmd: 'python3', ext: 'py' },
+        'ruby' => { cmd: 'ruby',    ext: 'rb' },
+        'node' => { cmd: 'node',    ext: 'js' },
+        'php' => { cmd: 'php', ext: 'php' }
+      }.freeze
+
+      def execute(code:, language:, stdin: '', timeout_seconds: DEFAULT_TIMEOUT)
+        spec = INTERPRETERS[language.to_s]
+        return failure("unsupported language: #{language}") unless spec
+
+        invoke(spec, language.to_s, code, stdin, timeout_seconds)
+      end
+
+      private
+
+      def invoke(spec, language, code, stdin, timeout_seconds)
+        Dir.mktmpdir('trmnlp-tx-') do |dir|
+          output_path = File.join(dir, 'output.json')
+          src_path = File.join(dir, "transform.#{spec[:ext]}")
+          File.write(src_path, Wrapper.for(language, code, sink_for(language, output_path)))
+          run_process(spec[:cmd], src_path, stdin, timeout_seconds, output_path)
+        end
+      rescue Errno::ENOENT, Errno::EACCES => e
+        failure("interpreter not available: #{e.message}")
+      end
+
+      def run_process(cmd, src_path, stdin, timeout_seconds, output_path)
+        started = monotonic_ms
+
+        Open3.popen3(cmd, src_path) do |stdin_io, stdout_io, stderr_io, wait_thr|
+          stdin_io.write(stdin)
+          stdin_io.close
+
+          unless wait_thr.join(timeout_seconds)
+            kill(wait_thr)
+            return failure("timeout after #{timeout_seconds}s", monotonic_ms - started)
+          end
+
+          build_result(stdout_io.read, stderr_io.read, wait_thr.value, output_path, monotonic_ms - started)
+        end
+      end
+
+      def build_result(stdout, stderr, status, output_path, duration_ms)
+        TransformClient::Result.new(
+          stdout: stdout,
+          stderr: stderr,
+          output: read_output(output_path),
+          exit_code: status.exitstatus || -1,
+          duration_ms: duration_ms,
+          error: nil
+        )
+      end
+
+      # A transform that crashes before writing leaves no output file; a
+      # permissions/IO error on read is treated the same — empty output,
+      # which the pipeline surfaces as a non-JSON-output failure.
+      def read_output(path)
+        File.exist?(path) ? File.read(path) : ''
+      rescue SystemCallError
+        ''
+      end
+
+      # TERM first, escalating to KILL only if the process outlives the
+      # grace period. join returns the moment it exits, so a process that
+      # dies promptly on TERM costs near-zero wait, not a fixed sleep.
+      def kill(wait_thr)
+        Process.kill('TERM', wait_thr.pid)
+        return if wait_thr.join(GRACE_PERIOD)
+
+        Process.kill('KILL', wait_thr.pid)
+      rescue Errno::ESRCH
+        # already exited between TERM and KILL — fine
+      end
+
+      def failure(message, duration_ms = 0)
+        TransformClient::Result.new(
+          stdout: '', stderr: '', output: '',
+          exit_code: -1, duration_ms: duration_ms, error: message
+        )
+      end
+
+      def monotonic_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) * 1000).to_i
+
+      # Language-specific tempfile sink snippets. The dispatch harness
+      # in Wrapper.* writes `output` (a serializable value); these
+      # snippets get it onto disk so the parent process can read it.
+      def sink_for(language, output_path)
+        case language
+        when 'python'
+          # A single statement — no block indentation to keep in sync with
+          # the Wrapper.python heredoc. open() is flushed and closed when
+          # CPython drops the temporary as the process exits.
+          "json.dump(output, open(#{output_path.inspect}, 'w'))"
+        when 'ruby'
+          "File.write(#{output_path.inspect}, JSON.generate(output))"
+        when 'node'
+          <<~JS.chomp
+            Promise.resolve(output)
+              .then((o) => require('fs').writeFileSync(#{output_path.inspect}, JSON.stringify(o)))
+              .catch((err) => { process.stderr.write(err.stack || String(err)); process.exit(1); });
+          JS
+        when 'php'
+          "file_put_contents(#{output_path.inspect}, json_encode($output));"
+        end
+      end
+    end
+  end
+end

data/lib/trmnlp/transform_backend/wrapper.rb

@@ -0,0 +1,113 @@
+# frozen_string_literal: true
+
+module TRMNLP
+  module TransformBackend
+    # Code wrappers shared by Subprocess and Http backends. Each method
+    # emits the canonical dispatch harness — read stdin → run user code
+    # → find run/transform/result/passthrough → serialize output — and
+    # delegates the final write to a language-appropriate `output_sink`
+    # snippet supplied by the caller. Subprocess writes to a tempfile
+    # path, Http writes to FD 3 for the production daemon to capture.
+    #
+    # Mirrors the hosted serverless runtime's code-wrapping behavior
+    # verbatim except for the configurable sink.
+    #
+    # NOTE: `output_sink` is spliced verbatim into the generated script as
+    # executable code. It MUST be trmnlp-generated (see Subprocess#sink_for
+    # and Http#sink_for) and never derived from user input or config — an
+    # attacker-influenced sink is arbitrary code execution in the transform
+    # process. Only `code` is untrusted; the sink is part of the harness.
+    module Wrapper
+      module_function
+
+      def python(code, output_sink)
+        <<~PYTHON
+          import sys, json, os
+          input = json.loads(sys.stdin.read())
+
+          #{code}
+
+          if callable(locals().get('run', None)):
+              output = run(input)
+          elif 'result' in dir():
+              output = result
+          else:
+              output = input
+          #{output_sink}
+        PYTHON
+      end
+
+      def ruby(code, output_sink)
+        <<~RUBY
+          require 'json'
+          input = JSON.parse($stdin.read)
+
+          #{code}
+
+          output = if defined?(run) == 'method'
+                     run(input)
+                   elsif defined?(result)
+                     result
+                   else
+                     input
+                   end
+          #{output_sink}
+        RUBY
+      end
+
+      # NOTE: node also accepts `function transform(input)` for
+      # production parity. Plugins authored against the hosted service
+      # using `transform` would otherwise silently pass input through.
+      def node(code, output_sink)
+        <<~JS
+          const input = JSON.parse(require('fs').readFileSync(0, 'utf8'));
+
+          #{code}
+
+          let output;
+          if (typeof run === "function") {
+            output = run(input);
+          } else if (typeof transform === "function") {
+            output = transform(input);
+          } else if (typeof result !== "undefined") {
+            output = result;
+          } else {
+            output = input;
+          }
+          #{output_sink}
+        JS
+      end
+
+      # NOTE: strips a leading `<?php` tag from user code so plugin
+      # authors can write the file as a standalone .php script. The
+      # hosted service does the same.
+      def php(code, output_sink)
+        cleaned = code.sub(/\A\s*<\?php\s*/, '')
+        <<~PHP
+          <?php
+          $input = json_decode(file_get_contents('php://stdin'), true);
+
+          #{cleaned}
+
+          if (function_exists('run')) {
+              $output = run($input);
+          } elseif (isset($result)) {
+              $output = $result;
+          } else {
+              $output = $input;
+          }
+          #{output_sink}
+        PHP
+      end
+
+      def for(language, code, output_sink)
+        case language.to_s
+        when 'python' then python(code, output_sink)
+        when 'ruby'   then ruby(code, output_sink)
+        when 'node'   then node(code, output_sink)
+        when 'php'    then php(code, output_sink)
+        end
+      end
+    end
+  end
+end