phantom_mcp

phoenix_tracker_available? =
  Code.ensure_loaded?(Phoenix.Tracker) and Code.ensure_loaded?(Phoenix.PubSub)

defmodule Phantom.Tracker do
  @moduledoc """
  Track open streams so that notifications and requests can be sent to clients.

  Add to your supervision tree:

      {Phoenix.PubSub, name: MyApp.PubSub},
      {Phantom.Tracker, [name: Phantom.Tracker, pubsub_server: MyApp.PubSub]},

  For example, a request may need to elicit more input from the client, so the
  first request stream will remain open, and the notification stream will send
  and new request to the client, and the client will POST its response. The
  new response connection will notify the first request connection with the result
  and the tool can continue with the elicited information.

  See `m:Phantom#module-persistent-streams` section for more information.
  """

  if phoenix_tracker_available? do
    Code.eval_quoted(quote(do: use(Phoenix.Tracker)), [], __ENV__)
  end

  @available phoenix_tracker_available?
  @sessions "phantom:sessions"
  @requests "phantom:requests"
  @resources "phantom:resources"

  def available?, do: @available
  def resource_subscription_topic, do: @resources
  def requests_topic, do: @requests
  def sessions_topic, do: @sessions

  @doc false
  if @available do
    def start_link(opts) do
      opts = Keyword.merge([name: __MODULE__], opts)
      Phoenix.Tracker.start_link(__MODULE__, opts, opts)
    end
  else
    def start_link(_opts), do: :error
  end

  @doc false
  if @available do
    def init(opts) do
      server = Keyword.fetch!(opts, :pubsub_server)
      {:ok, %{pubsub_server: server, node_name: Phoenix.PubSub.node_name(server)}}
    end
  else
    def init(_opts), do: :ignore
  end

  @doc "Track a request PID"
  if @available do
    def track_request(pid, request_id, meta \\ %{}) do
      Phoenix.Tracker.track(__MODULE__, pid, @requests, request_id, Map.put(meta, :pid, pid))
    rescue
      _ -> {:error, :tracker_not_in_supervision_tree}
    end
  else
    def track_request(_pid, _request_id, _meta \\ %{}), do: {:error, :not_available}
  end

  @doc "Track a session PID"
  if @available do
    require Logger

    def track_session(pid, session_id, meta \\ %{}) do
      Phoenix.Tracker.track(__MODULE__, pid, @sessions, session_id, Map.put(meta, :pid, pid))
    rescue
      _ ->
        if not warned?() do
          Logger.warning(
            "Phoenix.PubSub is available, but Phantom.Tracker is not in the supervision tree. Please add this to your supervision tree: `{Phantom.Tracker, [name: Phantom.Tracker, pubsub_server: MyApp.PubSub]}`"
          )
        end

        {:error, :tracker_not_in_supervision_tree}
    end

    defp warned? do
      :persistent_term.get({Phantom.Tracker, :warned}, false) == true ||
        :persistent_term.put({Phantom.Tracker, :warned}, true) != :ok
    end
  else
    def track_session(_pid, _session_id, _meta \\ %{}), do: {:error, :not_available}
  end

  @doc "Update metadata for a tracked session.

  Falls back to process dictionary when Phoenix.Tracker is not tracking
  (e.g. stdio transport).
  "
  if @available do
    def update_session_meta(session_id, meta) do
      case Phoenix.Tracker.get_by_key(__MODULE__, @sessions, session_id) do
        [{_key, %{pid: pid} = existing_meta} | _] ->
          Phoenix.Tracker.update(
            __MODULE__,
            pid,
            @sessions,
            session_id,
            Map.merge(existing_meta, meta)
          )

        _ ->
          update_local_session_meta(meta)
      end
    rescue
      _ -> update_local_session_meta(meta)
    end
  else
    def update_session_meta(_session_id, meta), do: update_local_session_meta(meta)
  end

  @doc "Get metadata for a tracked session.

  Falls back to process dictionary when Phoenix.Tracker is not tracking.
  "
  if @available do
    def get_session_meta(session_id) do
      case Phoenix.Tracker.get_by_key(__MODULE__, @sessions, session_id) do
        [{_key, meta} | _] -> meta
        _ -> Process.get(:phantom_session_meta, %{})
      end
    rescue
      _ -> Process.get(:phantom_session_meta, %{})
    end
  else
    def get_session_meta(_session_id), do: Process.get(:phantom_session_meta, %{})
  end

  defp update_local_session_meta(meta) do
    existing = Process.get(:phantom_session_meta, %{})
    Process.put(:phantom_session_meta, Map.merge(existing, meta))
    :ok
  end

  @doc "Return a list of all open sessions"
  if @available do
    def list_sessions do
      Phoenix.Tracker.list(__MODULE__, @sessions)
    rescue
      _ -> []
    end
  else
    def list_sessions, do: []
  end

  @doc "Return a list of all open requests"
  if @available do
    def list_requests do
      Phoenix.Tracker.list(__MODULE__, @requests)
    rescue
      _ -> []
    end
  else
    def list_requests, do: []
  end

  @doc "Return a list of all listening for resources"
  if @available do
    def list_resource_listeners do
      Phoenix.Tracker.list(__MODULE__, @resources)
    rescue
      _ -> []
    end
  else
    def list_resource_listeners, do: []
  end

  @doc "Fetch the PID of the open request by ID"
  def get_request(%Phantom.Request{id: request_id}), do: get_request(request_id)

  if @available do
    def get_request(request_id) do
      case Phoenix.Tracker.get_by_key(__MODULE__, @requests, request_id) do
        [{_key, %{pid: pid}} | _] ->
          if Process.alive?(pid) do
            pid
          else
            Phoenix.Tracker.untrack(__MODULE__, pid)
            nil
          end

        _ ->
          nil
      end
    rescue
      _ -> nil
    end
  else
    def get_request(_request_id), do: nil
  end

  @doc "Fetch the PID of the open session by ID"
  def get_session(%Phantom.Session{id: session_id}), do: get_session(session_id)

  if @available do
    def get_session(session_id) do
      case Phoenix.Tracker.get_by_key(__MODULE__, @sessions, session_id) do
        [{_key, %{pid: pid}} | _] ->
          if Process.alive?(pid) do
            pid
          else
            Phoenix.Tracker.untrack(__MODULE__, pid)
            nil
          end

        _ ->
          nil
      end
    rescue
      _ -> nil
    end
  else
    def get_session(_session_id), do: nil
  end

  @doc "Untrack the processe for everything"
  if @available do
    def untrack(pid), do: Phoenix.Tracker.untrack(__MODULE__, pid)
  else
    def untrack(_pid), do: :ok
  end

  @doc "Untrack any processes for the session"
  def untrack_session(%Phantom.Session{id: session_id}), do: untrack_session(session_id)

  if @available do
    def untrack_session(session_id) do
      tracked =
        try do
          Phoenix.Tracker.get_by_key(__MODULE__, @sessions, session_id)
        rescue
          _ -> []
        end

      Enum.each(tracked, fn {_key, %{pid: pid}} -> Phoenix.Tracker.untrack(__MODULE__, pid) end)

      :ok
    rescue
      _ -> :ok
    end
  else
    def untrack_session(_session_id), do: :ok
  end

  @doc "Untrack any processes for the request"
  def untrack_request(%Phantom.Request{id: request_id}), do: untrack_request(request_id)

  if @available do
    def untrack_request(request_id) do
      tracked =
        try do
          Phoenix.Tracker.get_by_key(__MODULE__, @requests, request_id)
        rescue
          _ -> []
        end

      Enum.each(tracked, fn {_key, %{pid: pid}} -> Phoenix.Tracker.untrack(__MODULE__, pid) end)

      :ok
    rescue
      _ -> :ok
    end
  else
    def untrack_request(_request_id), do: :ok
  end

  @doc "Subscribe the process to resource notifications from the PubSub on topic #{inspect(@resources)}"
  if @available do
    def subscribe_resource(uri) do
      Phoenix.Tracker.track(__MODULE__, self(), @resources, uri, %{pid: self()})
    rescue
      _ -> {:error, :tracker_not_in_supervision_tree}
    end
  else
    def subscribe_resource(_uri), do: {:error, :not_available}
  end

  @doc "Unsubscribe the process to resource notifications from the PubSub on topic #{inspect(@resources)}"
  if @available do
    def unsubscribe_resource(uri) do
      Phoenix.Tracker.untrack(__MODULE__, self(), @resources, uri)
    rescue
      _ -> :ok
    end
  else
    def unsubscribe_resource(_uri), do: {:error, :not_available}
  end

  @doc "Notify any listening MCP sessions that the resource has updated"
  if @available do
    def notify_resource_updated(uri) do
      tracked =
        try do
          Phoenix.Tracker.get_by_key(__MODULE__, @resources, uri)
        rescue
          _ -> []
        end

      {:ok,
       Enum.count(tracked, fn {_key, %{pid: pid}} ->
         GenServer.cast(pid, {:resource_updated, uri})
       end)}
    rescue
      _ -> {:error, :tracker_not_in_supervision_tree}
    end
  else
    def notify_resource_updated(_), do: {:ok, 0}
  end

  @doc "Notify any listening MCP sessions that the list of tools has updated"
  if @available do
    def notify_tool_list do
      {:ok,
       Enum.count(list_sessions(), fn {session_id, _} ->
         if pid = get_session(session_id), do: GenServer.cast(pid, :tools_updated)
       end)}
    end
  else
    def notify_tool_list, do: {:ok, 0}
  end

  @doc "Notify any listening MCP sessions that the list of prompts has updated"
  if @available do
    def notify_prompt_list do
      {:ok,
       Enum.count(list_sessions(), fn {session_id, _} ->
         if pid = get_session(session_id), do: GenServer.cast(pid, :prompts_updated)
       end)}
    end
  else
    def notify_prompt_list, do: {:ok, 0}
  end

  @doc "Notify any listening MCP sessions that the list of prompts has updated"
  if @available do
    def notify_resource_list do
      {:ok,
       Enum.count(list_sessions(), fn {session_id, _} ->
         if pid = get_session(session_id), do: GenServer.cast(pid, :resources_updated)
       end)}
    end
  else
    def notify_resource_list, do: {:ok, 0}
  end

  @doc "Notify tracked sessions that a URL elicitation has completed"
  if @available do
    def notify_elicitation_complete(elicitation_id) do
      tracked =
        try do
          Phoenix.Tracker.get_by_key(__MODULE__, @requests, elicitation_id)
        rescue
          _ -> []
        end

      Enum.each(tracked, fn {_key, %{type: :elicitation, pid: pid}} ->
        GenServer.cast(pid, {:send, Phantom.Request.elicitation_complete(elicitation_id)})
        Phoenix.Tracker.untrack(__MODULE__, pid, @requests, elicitation_id)
      end)

      {:ok, length(tracked)}
    end
  else
    def notify_elicitation_complete(_elicitation_id), do: {:ok, 0}
  end

  @doc false
  if @available do
    def handle_diff(diff, state) do
      for {topic, {joins, leaves}} <- diff do
        for {key, meta} <- joins do
          msg = {:join, key, meta}
          Phoenix.PubSub.direct_broadcast!(state.node_name, state.pubsub_server, topic, msg)
        end

        for {key, meta} <- leaves do
          msg = {:leave, key, meta}
          Phoenix.PubSub.direct_broadcast!(state.node_name, state.pubsub_server, topic, msg)
        end
      end

      {:ok, state}
    end
  else
    def handle_diff(_diff, state), do: {:ok, state}
  end
end