lib/omni/tools/repl/extension.ex

defmodule Omni.Tools.Repl.Extension do
  @moduledoc """
  Behaviour and struct for extending the REPL sandbox environment.

  Extensions inject code and documentation into the sandbox. There are
  two ways to define an extension:

  ## Inline extensions

  Use `new/1` to create lightweight extensions without defining a module.
  At least one of `:code` or `:description` must be provided.

      # Description-only — document available host modules
      Extension.new(description: "Req and Jason are available.")

      # Code-only — inject setup code
      Extension.new(code: quote(do: defmodule(Helper, do: def(ping, do: :pong))))

      # Both
      Extension.new(
        code: quote(do: defmodule(Helper, do: def(ping, do: :pong))),
        description: "- `Helper.ping/0` — returns `:pong`"
      )

  ## Module-based extensions

  For reusable extensions, define a module implementing the behaviour.
  Both callbacks are required.

      defmodule MyApp.ReplExtension do
        @behaviour Omni.Tools.Repl.Extension

        @impl true
        def code(opts) do
          api_key = Keyword.fetch!(opts, :api_key)

          quote do
            defmodule MyAPI do
              def fetch(path) do
                Req.get!(path, headers: [{"authorization", unquote(api_key)}]).body
              end
            end
          end
        end

        @impl true
        def description(_opts) do
          \"""
          ## MyAPI
          - `MyAPI.fetch(path)` — authenticated GET request
          \"""
        end
      end

  ## Usage

  Pass extensions to `Omni.Tools.Repl.new/1`:

      Omni.Tools.Repl.new(
        extensions: [
          {MyApp.ReplExtension, api_key: "sk-..."},
          Extension.new(description: "Extra context for the model")
        ]
      )
  """

  @typedoc "Code to evaluate in the sandbox — AST (preferred) or a string."
  @type setup_code :: String.t() | Macro.t()

  @typedoc "A resolved extension with optional code and description."
  @type t :: %__MODULE__{
          code: setup_code() | nil,
          description: String.t() | nil
        }

  defstruct [:code, :description]

  @doc """
  Returns code to evaluate in the sandbox before the user's code.

  Receives the opts from the `{module, opts}` tuple in the extensions
  list. Return a quoted expression (preferred) or a code string. The
  code is evaluated in the peer node before IO capture begins.
  """
  @callback code(opts :: keyword()) :: setup_code()

  @doc """
  Returns a description fragment appended to the REPL tool description.

  Receives the same opts as `code/1`. The returned string should document
  the APIs made available by `code/1` so the agent knows how to use them.
  """
  @callback description(opts :: keyword()) :: String.t()

  @doc """
  Creates an inline extension.

  At least one of `:code` or `:description` must be provided.

      Extension.new(description: "Req is available. Do not Mix.install it.")
      Extension.new(code: "defmodule(H, do: def(hi, do: :hello))")
  """
  @spec new(keyword()) :: t()
  def new(opts \\ []) do
    code = Keyword.get(opts, :code)
    description = Keyword.get(opts, :description)

    if is_nil(code) and is_nil(description) do
      raise ArgumentError, "extension requires at least one of :code or :description"
    end

    %__MODULE__{code: code, description: description}
  end
end