From 86a6d1f041483c27834613eb22871ca5fbbf2143 Mon Sep 17 00:00:00 2001 From: Monte Brown Date: Sat, 24 Jan 2026 18:53:49 -0500 Subject: [PATCH] StatusMonitor recovery from process restarts When a client process restarts (e.g., MCP server goes down, supervisor restarts the client), StatusMonitor previously reported :process_dead permanently because it stored a static PID at registration time. Now register_client_by_name/3 stores the process name (atom) instead of resolving to a PID, and get_client_status/1 and health_check/1 resolve dynamically via Process.whereis/1. This allows automatic recovery when supervisors restart clients under the same name. Existing register_client/2 behavior (PID-based) unchanged for backward compatibility. --- lib/langchain_mcp/status_monitor.ex | 31 ++++++++- test/langchain_mcp/status_monitor_test.exs | 77 +++++++++++++++++++++- 2 files changed, 105 insertions(+), 3 deletions(-) diff --git a/lib/langchain_mcp/status_monitor.ex b/lib/langchain_mcp/status_monitor.ex index 4eab32c..d3e2404 100644 --- a/lib/langchain_mcp/status_monitor.ex +++ b/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 diff --git a/test/langchain_mcp/status_monitor_test.exs b/test/langchain_mcp/status_monitor_test.exs index b630900..783a72e 100644 --- a/test/langchain_mcp/status_monitor_test.exs +++ b/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,