Implement plural mcp functionality (#1408)

Author: michaeljguarino

apps/api/lib/api/application.ex

@@ -8,6 +8,7 @@ defmodule Api.Application do
     ApiWeb.Plugs.MetricsExporter.setup()
     children = [
       ApiWeb.Endpoint,
+      {Bandit, plug: Core.MCP.Router, port: 3000},
       {FT.K8S.TrafficDrainHandler, Core.drain_config()},
       {Cluster.Supervisor, [topologies, [name: Api.ClusterSupervisor]]},
     ]

apps/core/lib/core/mcp/jwt.ex

@@ -0,0 +1,24 @@
+defmodule Core.MCP.Jwt do
+  use Joken.Config
+  use Nebulex.Caching
+
+  @alg "ES256"
+  @ttl :timer.minutes(15)
+
+  def exchange(jwt) do
+    with {:ok, %{"keys" => [jws | _]}} <- get_jwks(),
+     do: verify_and_validate(jwt, signer(jws))
+  end
+
+  def signer(jws), do: Joken.Signer.create(@alg, jws)
+
+  @decorate cacheable(cache: Core.Cache, key: :console_jws, opts: [ttl: @ttl])
+  defp get_jwks() do
+    Path.join([Core.conf(:console_url), "mcp", ".well-known", "jwks.json"])
+    |> HTTPoison.get()
+    |> case do
+      {:ok, %HTTPoison.Response{status_code: 200, body: body}} -> Jason.decode(body)
+      _ -> {:error, "failed to fetch jwks"}
+    end
+  end
+end

apps/core/lib/core/mcp/router.ex

@@ -0,0 +1,51 @@
+defmodule Core.MCP.Router do
+  use Plug.Router
+  alias Plug.Conn
+  alias Core.MCP.Jwt
+
+  plug Plug.Parsers,
+   parsers: [:urlencoded, :json],
+   pass: ["text/*"],
+   json_decoder: Jason
+
+  plug :match
+  plug :ensure_session_id
+  plug :authorize
+  plug :dispatch
+
+  forward "/sse", to: SSE.ConnectionPlug
+  forward "/message", to: SSE.ConnectionPlug
+
+  match _ do
+   send_resp(conn, 404, "Not found")
+  end
+
+  # Middleware to ensure session ID exists
+  def ensure_session_id(%Conn{query_params: %{"sessionId" => id}} = conn, _) when is_binary(id),
+    do: conn
+  def ensure_session_id(conn, _opts) do
+    session_id = Base.encode16(:crypto.strong_rand_bytes(8), case: :lower)
+    put_in(conn.query_params["sessionId"], session_id)
+  end
+
+  def authorize(conn, _) do
+    with ["Bearer " <> token | _] <- get_req_header(conn, "authorization"),
+         true <- authorized?(token) do
+      conn
+    else
+      _ ->
+        send_resp(conn, 403, "unauthorized")
+        |> halt()
+    end
+  end
+
+  @groups ~w(sales ops)
+
+  defp authorized?(token) do
+    case Jwt.exchange(token) do
+      {:ok, %{"admin" => true}} -> true
+      {:ok, %{"groups" => groups}} -> Enum.any?(groups, & &1 in @groups)
+      _ -> false
+    end
+  end
+end

apps/core/lib/core/mcp/server.ex

@@ -0,0 +1,79 @@
+defmodule Core.MCP.Server do
+  use MCPServer
+  alias Core.MCP.Tools.{
+    Account,
+    Enterprise,
+    CloudConsole,
+    RemoveEnterprise
+  }
+
+  @tools [Account, Enterprise, CloudConsole, RemoveEnterprise]
+  @by_name Map.new(@tools, & {&1.name(), &1})
+
+  @protocol_version "2024-11-05"
+
+  @impl true
+  def handle_ping(request_id) do
+    {:ok, %{jsonrpc: "2.0", id: request_id, result: %{}}}
+  end
+
+  @impl true
+  def handle_initialize(request_id, params) do
+    case validate_protocol_version(params["protocolVersion"]) do
+      :ok ->
+        {:ok,
+         %{
+           jsonrpc: "2.0",
+           id: request_id,
+           result: %{
+             protocolVersion: @protocol_version,
+             capabilities: %{
+               tools: %{
+                 listChanged: true
+               }
+             },
+             serverInfo: %{
+               name: "Plural MCP Server",
+               version: "0.1.0"
+             }
+           }
+         }}
+
+      {:error, reason} ->
+        {:error, reason}
+    end
+  end
+
+  @impl true
+  def handle_list_tools(request_id, _params) do
+    tools = Enum.map(@tools, & %{
+      name: &1.name(),
+      description: &1.description(),
+      inputSchema: &1.schema()
+    })
+    {:ok, %{jsonrpc: "2.0", id: request_id, result: %{tools: tools}}}
+  end
+
+  @impl true
+  def handle_call_tool(request_id, %{"name" => name} = params) do
+    with {:tool, t} when not is_nil(t) <- {:tool, @by_name[name]},
+         {:ok, result} <- t.invoke(params["arguments"]) do
+       {:ok, %{jsonrpc: "2.0", id: request_id, result: %{content: [%{type: "text", text: result}]}}}
+    else
+      {:tool, _} ->
+        {:error, %{
+          jsonrpc: "2.0",
+          id: request_id,
+          error: %{
+            code: -32601,
+            message: "Method not found",
+            data: %{
+              name: "no tool with name #{name}"
+            }
+          }
+        }}
+      {:error, err} ->
+        {:ok, %{jsonrpc: "2.0", id: request_id, result: %{isError: true, content: [%{type: "text", text: err}]}}}
+    end
+  end
+end

apps/core/lib/core/mcp/tool.ex

@@ -0,0 +1,21 @@
+defmodule Core.MCP.Tool do
+  @doc """
+  The mcp server tool description
+  """
+  @callback description() :: binary
+
+  @doc """
+  jsonapi format input schema
+  """
+  @callback schema() :: map
+
+  @doc """
+  name of the tool
+  """
+  @callback name() :: binary
+
+  @doc """
+  actually calls the tool
+  """
+  @callback invoke(args :: map) :: {:ok, binary} | {:error, binary}
+end

apps/core/lib/core/mcp/tools/account.ex

@@ -0,0 +1,33 @@
+defmodule Core.MCP.Tools.Account do
+  @behaviour Core.MCP.Tool
+  alias Core.Repo
+  alias Core.Services.Users
+  alias Core.Schema.User
+
+  def name(), do: "account"
+
+  def description(), do: "Fetches the full Plural account given a user email, also including subscription and plan details"
+
+  def schema() do
+    %{
+      type: "object",
+      required: ["email"],
+      properties: %{
+        email: %{
+          type: "string",
+          description: "The email of the user whose account you want to look up"
+        }
+      }
+    }
+  end
+
+  def invoke(%{"email" => email}) do
+    with %User{} = user <- Users.get_user_by_email(email) do
+      user = Repo.preload(user, [account: [subscription: :plan]])
+      Jason.encode(user)
+    else
+      _ -> {:ok, "no user with email #{email}"}
+    end
+  end
+  def invoke(_), do: {:error, "email is required"}
+end

apps/core/lib/core/mcp/tools/cloud_console.ex

@@ -0,0 +1,37 @@
+defmodule Core.MCP.Tools.CloudConsole do
+  @behaviour Core.MCP.Tool
+  alias Core.Repo
+  alias Core.Services.Cloud
+  alias Core.Schema.{ConsoleInstance}
+
+  def name(), do: "console_cluster"
+
+  def description(), do: "Fetches the details of a cloud console"
+
+  def schema() do
+    %{
+      type: "object",
+      required: ["name"],
+      properties: %{
+        name: %{
+          type: "string",
+          description: "The name of the console instance"
+        }
+      }
+    }
+  end
+
+  def invoke(%{"name" => name}) do
+    with %ConsoleInstance{} = console <- Cloud.get_instance_by_name(name),
+         console = Repo.preload(console, [:cluster, owner: [account: [subscription: :plan]]]) do
+      Jason.encode(format(console))
+    else
+      _ -> {:ok, "no instance with name #{name}"}
+    end
+  end
+  def invoke(_), do: {:error, "email is required"}
+
+  defp format(console) do
+    Map.take(console, ~w(id type name status subdomain url cloud size region cluster owner)a)
+  end
+end

apps/core/lib/core/mcp/tools/enterprise.ex

@@ -0,0 +1,33 @@
+defmodule Core.MCP.Tools.Enterprise do
+  @behaviour Core.MCP.Tool
+  alias Core.Services.{Accounts, Payments}
+  alias Core.Schema.Account
+
+  def name(), do: "add_enterprise_plan"
+
+  def description(), do: "Adds an account to the main enterprise plan"
+
+  def schema() do
+    %{
+      type: "object",
+      required: ["account_id"],
+      properties: %{
+        account_id: %{
+          type: "string",
+          description: "The account id to add to the enterprise plan.  You can fetch the id by first fetching the users account info."
+        }
+      }
+    }
+  end
+
+  def invoke(%{"account_id" => id}) do
+    with %Account{} = account <- Accounts.get_account(id),
+         {:ok, _} <- Payments.remove_trial(account),
+         {:ok, _} <- Payments.setup_enterprise_plan(account.id) do
+      {:ok, "added the account to the enterprise plan successfully"}
+    else
+      _ -> {:ok, "no account with id #{id}"}
+    end
+  end
+  def invoke(_), do: {:error, "account id is required"}
+end

apps/core/lib/core/mcp/tools/remove_enterprise.ex

@@ -0,0 +1,34 @@
+defmodule Core.MCP.Tools.RemoveEnterprise do
+  @behaviour Core.MCP.Tool
+  alias Core.Services.{Accounts, Payments}
+  alias Core.Schema.Account
+
+  def name(), do: "remove_enterprise_plan"
+
+  def description(), do: "removes an account to the main enterprise plan"
+
+  def schema() do
+    %{
+      type: "object",
+      required: ["account_id"],
+      properties: %{
+        account_id: %{
+          type: "string",
+          description: "The account id to add to the enterprise plan. You can fetch the id by first querying a users account"
+        }
+      }
+    }
+  end
+
+  def invoke(%{"account_id" => id}) do
+    with {:account, %Account{} = account} <- {:account, Accounts.get_account(id)},
+         {:ok, _} <- Payments.remove_enterprise_plan(account.id) do
+      {:ok, "removed the account from the enterprise plan successfully"}
+    else
+      {:account, _} -> {:ok, "no account with id #{id}"}
+      err -> err
+    end
+  end
+
+  def invoke(_), do: {:error, "account id is required"}
+end

apps/core/lib/core/services/accounts.ex

@@ -32,6 +32,8 @@ defmodule Core.Services.Accounts do
 
   def get_account!(id), do: Core.Repo.get!(Account, id)
 
+  def get_account(id), do: Core.Repo.get(Account, id)
+
   def get_group!(id), do: Core.Repo.get!(Group, id)
 
   def get_group_by_name(aid, name), do: Core.Repo.get_by(Group, name: name, account_id: aid)

apps/core/lib/core/services/payments.ex

@@ -675,6 +675,7 @@ defmodule Core.Services.Payments do
   defp discount(amount, :yearly), do: round(9 * amount / 10) * 12
   defp discount(amount, _), do: amount
 
+  @spec setup_enterprise_plan(User.t | binary) :: platform_sub_resp
   def setup_enterprise_plan(%User{} = user) do
     case Repo.preload(user, [account: [subscription: :plan]]) do
       %{account: %Account{} = account} ->
@@ -683,15 +684,25 @@ defmodule Core.Services.Payments do
       _ -> setup_enterprise_plan(user.account_id)
     end
   end
-
-  def setup_enterprise_plan(account_id) do
+  def setup_enterprise_plan(account_id) when is_binary(account_id) do
     plan = get_platform_plan_by_name!("Enterprise")
 
     %PlatformSubscription{account_id: account_id}
     |> PlatformSubscription.changeset(%{plan_id: plan.id})
     |> Core.Repo.insert()
   end
 
+  @spec remove_enterprise_plan(binary) :: platform_sub_resp
+  def remove_enterprise_plan(account_id) do
+    Repo.get_by(PlatformSubscription, account_id: account_id)
+    |> Repo.preload([:plan])
+    |> case do
+      %PlatformSubscription{plan: %PlatformPlan{enterprise: true}} = sub ->
+        Repo.delete(sub)
+      _ -> {:error, "account #{account_id} is not on an enterprise plan"}
+    end
+  end
+
   @doc """
   Creates a plan and associated add-ons.  Transactionally creates them in stripe, then
   restitches them back into the db.

apps/core/mix.exs

@@ -76,6 +76,7 @@ defmodule Core.MixProject do
       {:bamboo, "~> 2.0"},
       {:parallax, "~> 1.0"},
       {:bourne, "~> 1.1"},
+      {:mcp_sse, "~> 0.1.0"},
       {:flow, "~> 0.15.0"},
       {:joken, "~> 2.5.0"},
       {:guardian, "~> 1.2.1"},