StatusMonitor recovery from process restarts

lib/langchain_mcp/status_monitor.ex

@@ -83,6 +83,15 @@ defmodule LangChain.MCP.StatusMonitor do
   """
   def get_client_status(name) when is_atom(name) do
     case Registry.lookup(@registry_name, name) do
+      [{_pid, %{process_name: process_name}}] ->
+        case Process.whereis(process_name) do
+          pid when is_pid(pid) ->
+            {:ok, %{pid: pid}}
+
+          nil ->
+            {:error, {:client_unavailable, :process_dead}}
+        end
+
       [{_pid, %{pid: client_pid}}] ->
         if Process.alive?(client_pid) do
           {:ok, %{pid: client_pid}}
@@ -125,6 +134,17 @@ defmodule LangChain.MCP.StatusMonitor do
   """
   def health_check(name) when is_atom(name) do
     case Registry.lookup(@registry_name, name) do
+      [{_pid, %{process_name: process_name}}] ->
+        ping_result =
+          case Process.whereis(process_name) do
+            pid when is_pid(pid) -> :pong
+            nil -> {:error, :process_dead}
+          end
+
+        status_result = get_client_status(name)
+
+        {ping_result, status_result}
+
       [{_pid, %{pid: client_pid}}] ->
         ping_result =
           if Process.alive?(client_pid) do
@@ -282,7 +302,16 @@ defmodule LangChain.MCP.StatusMonitor do
   defp wait_and_register(module, registry_name, deadline) do
     case Process.whereis(module) do
       pid when is_pid(pid) ->
-        register_client(registry_name, pid)
+        # Store the process name for dynamic resolution, not the PID
+        case Registry.register(@registry_name, registry_name, %{process_name: module}) do
+          {:ok, _pid} ->
+            Logger.info("Registered MCP client #{registry_name} (#{module}) for status monitoring")
+            {:ok, :registered}
+
+          {:error, {:already_registered, _}} ->
+            Logger.debug("MCP client #{registry_name} already registered")
+            {:ok, :already_registered}
+        end
 
       nil ->
         if System.monotonic_time(:millisecond) < deadline do

test/langchain_mcp/status_monitor_test.exs

@@ -405,9 +405,11 @@ defmodule LangChain.MCP.StatusMonitorTest do
     end
 
     test "returns already_registered if client was already registered" do
-      # Start and register a process
+      # Start and register a process by name first
       {:ok, pid} = Agent.start_link(fn -> %{} end, name: AlreadyRegisteredModule)
-      StatusMonitor.register_client(:already_reg, pid)
+
+      assert {:ok, :registered} =
+               StatusMonitor.register_client_by_name(AlreadyRegisteredModule, :already_reg)
 
       # Try to register again using module name
       assert {:ok, :already_registered} =
@@ -471,6 +473,77 @@ defmodule LangChain.MCP.StatusMonitorTest do
       Agent.stop(pid3)
     end
 
+    test "stores process name in Registry instead of PID" do
+      {:ok, pid} = Agent.start_link(fn -> %{} end, name: StorageCheckModule)
+
+      assert {:ok, :registered} =
+               StatusMonitor.register_client_by_name(StorageCheckModule, :storage_check)
+
+      # Registry value should contain process_name, not pid
+      [{_reg_pid, value}] = Registry.lookup(:langchain_mcp_clients, :storage_check)
+      assert %{process_name: StorageCheckModule} = value
+      refute Map.has_key?(value, :pid)
+
+      Agent.stop(pid)
+    end
+
+    test "recovers from process restart with same name" do
+      Process.flag(:trap_exit, true)
+
+      # Start a named process and register it
+      {:ok, pid1} = Agent.start_link(fn -> %{} end, name: RestartableModule)
+
+      assert {:ok, :registered} =
+               StatusMonitor.register_client_by_name(RestartableModule, :restartable)
+
+      # Verify status is healthy with original PID
+      assert {:ok, %{pid: ^pid1}} = StatusMonitor.get_client_status(:restartable)
+
+      # Kill the process
+      Process.exit(pid1, :kill)
+      Process.sleep(10)
+
+      # Status should report process_dead while no process is registered under the name
+      assert {:error, {:client_unavailable, :process_dead}} =
+               StatusMonitor.get_client_status(:restartable)
+
+      # Start a new process with the same name (simulating supervisor restart)
+      {:ok, pid2} = Agent.start_link(fn -> %{} end, name: RestartableModule)
+      assert pid1 != pid2
+
+      # Status should now report healthy with the NEW PID
+      assert {:ok, %{pid: ^pid2}} = StatusMonitor.get_client_status(:restartable)
+
+      # Health check should also recover
+      assert {:pong, {:ok, %{pid: ^pid2}}} = StatusMonitor.health_check(:restartable)
+
+      Agent.stop(pid2)
+    end
+
+    test "recovery works with dashboard_status" do
+      Process.flag(:trap_exit, true)
+
+      {:ok, pid1} = Agent.start_link(fn -> %{} end, name: DashboardRecoveryModule)
+
+      assert {:ok, :registered} =
+               StatusMonitor.register_client_by_name(DashboardRecoveryModule, :dashboard_recovery)
+
+      # Kill and restart
+      Process.exit(pid1, :kill)
+      Process.sleep(10)
+
+      {:ok, pid2} = Agent.start_link(fn -> %{} end, name: DashboardRecoveryModule)
+
+      # Dashboard should show healthy with new PID
+      result = StatusMonitor.dashboard_status()
+      client_status = result.clients[:dashboard_recovery]
+      assert client_status.status == :healthy
+      assert client_status.alive? == true
+      assert client_status.pid == pid2
+
+      Agent.stop(pid2)
+    end
+
     test "works with LangChain.MCP.Client modules" do
       defmodule TestMCPClientForRegistration do
         use LangChain.MCP.Client,