Save current state before fixing MCP tool registration and proxying

2025-09-03 03:01:58 -07:00
parent 004ba6ab55
commit 4b7c4b6314
6 changed files with 167 additions and 45 deletions
--- a/.github/instructions/no-duplicate-files.instructions.md
+++ b/.github/instructions/no-duplicate-files.instructions.md
@@ -36,7 +36,7 @@ When you create duplicate files:

 ## The Human Is Right

-The human specifically said: "I fucking hate it when you do this retarded shit and recreate the same file with some adjective/verb but leave the original"
+The human specifically said: "Do not re-create the same file with some adjective/verb attached while leaving the original, instead, update the code and make it better, changes are good."

 **Listen to them.** They prefer file replacement over duplicates.

--- a/lib/agent_coordinator.ex
+++ b/lib/agent_coordinator.ex
@@ -255,8 +255,6 @@ defmodule AgentCoordinator do
    Map.get(server_status, :active_servers, 0)
  end

-  defp count_active_servers(_), do: 0
-
  defp get_server_status do
    # This would call UnifiedMCPServer to get external server status
    # For now, return a placeholder
--- a/lib/agent_coordinator/client.ex
+++ b/lib/agent_coordinator/client.ex
@@ -20,7 +20,7 @@ defmodule AgentCoordinator.Client do
  """

  use GenServer
-  alias AgentCoordinator.{EnhancedMCPServer, AutoHeartbeat}
+  alias AgentCoordinator.AutoHeartbeat

  defstruct [
    :agent_id,
@@ -108,11 +108,11 @@ defmodule AgentCoordinator.Client do
  # Server callbacks

  def init(config) do
-    # Register with enhanced MCP server
-    case EnhancedMCPServer.register_agent_with_session(
+    # Register with task registry
+    case AgentCoordinator.TaskRegistry.register_agent(
           config.agent_name,
           config.capabilities,
-           self()
+           session_pid: self()
         ) do
      {:ok, agent_id} ->
        state = %__MODULE__{
@@ -186,9 +186,9 @@ defmodule AgentCoordinator.Client do
  end

  def handle_call(:get_task_board, _from, state) do
-    case EnhancedMCPServer.get_enhanced_task_board() do
-      {:ok, board} ->
-        {:reply, {:ok, board}, update_last_heartbeat(state)}
+    case AgentCoordinator.TaskRegistry.get_task_board() do
+      task_board when is_map(task_board) ->
+        {:reply, {:ok, task_board}, update_last_heartbeat(state)}

      {:error, reason} ->
        {:reply, {:error, reason}, state}
@@ -270,12 +270,10 @@ defmodule AgentCoordinator.Client do
  # Private helpers

  defp enhanced_mcp_call(request, state) do
-    session_info = %{
-      agent_id: state.agent_id,
-      session_pid: state.session_pid
-    }
+    # Add agent_id to the request for the MCP server
+    request_with_agent = Map.put(request, "agent_id", state.agent_id)

-    case EnhancedMCPServer.handle_enhanced_mcp_request(request, session_info) do
+    case AgentCoordinator.MCPServer.handle_mcp_request(request_with_agent) do
      %{"result" => %{"content" => [%{"text" => response_json}]}} = response ->
        case Jason.decode(response_json) do
          {:ok, data} ->
@@ -304,7 +302,7 @@ defmodule AgentCoordinator.Client do
      }
    }

-    case EnhancedMCPServer.handle_enhanced_mcp_request(request) do
+    case AgentCoordinator.MCPServer.handle_mcp_request(request) do
      %{"result" => _} -> :ok
      %{"error" => %{"message" => message}} -> {:error, message}
      _ -> {:error, :unknown_heartbeat_error}
--- a/lib/agent_coordinator/mcp_server.ex
+++ b/lib/agent_coordinator/mcp_server.ex
@@ -51,13 +51,14 @@ defmodule AgentCoordinator.MCPServer do
      "inputSchema" => %{
        "type" => "object",
        "properties" => %{
+          "agent_id" => %{"type" => "string"},
          "id" => %{"type" => "string"},
          "name" => %{"type" => "string"},
          "workspace_path" => %{"type" => "string"},
          "description" => %{"type" => "string"},
          "metadata" => %{"type" => "object"}
        },
-        "required" => ["name", "workspace_path"]
+        "required" => ["agent_id", "name", "workspace_path"]
      }
    },
    %{
@@ -66,6 +67,7 @@ defmodule AgentCoordinator.MCPServer do
      "inputSchema" => %{
        "type" => "object",
        "properties" => %{
+          "agent_id" => %{"type" => "string"},
          "title" => %{"type" => "string"},
          "description" => %{"type" => "string"},
          "priority" => %{"type" => "string", "enum" => ["low", "normal", "high", "urgent"]},
@@ -86,7 +88,7 @@ defmodule AgentCoordinator.MCPServer do
            }
          }
        },
-        "required" => ["title", "description"]
+        "required" => ["agent_id", "title", "description"]
      }
    },
    %{
@@ -95,6 +97,7 @@ defmodule AgentCoordinator.MCPServer do
      "inputSchema" => %{
        "type" => "object",
        "properties" => %{
+          "agent_id" => %{"type" => "string"},
          "title" => %{"type" => "string"},
          "description" => %{"type" => "string"},
          "primary_codebase_id" => %{"type" => "string"},
@@ -107,7 +110,7 @@ defmodule AgentCoordinator.MCPServer do
            "enum" => ["sequential", "parallel", "leader_follower"]
          }
        },
-        "required" => ["title", "description", "primary_codebase_id", "affected_codebases"]
+        "required" => ["agent_id", "title", "description", "primary_codebase_id", "affected_codebases"]
      }
    },
    %{
@@ -138,8 +141,10 @@ defmodule AgentCoordinator.MCPServer do
      "inputSchema" => %{
        "type" => "object",
        "properties" => %{
+          "agent_id" => %{"type" => "string"},
          "codebase_id" => %{"type" => "string"}
-        }
+        },
+        "required" => ["agent_id"]
      }
    },
    %{
@@ -148,9 +153,10 @@ defmodule AgentCoordinator.MCPServer do
      "inputSchema" => %{
        "type" => "object",
        "properties" => %{
+          "agent_id" => %{"type" => "string"},
          "codebase_id" => %{"type" => "string"}
        },
-        "required" => ["codebase_id"]
+        "required" => ["agent_id", "codebase_id"]
      }
    },
    %{
@@ -158,7 +164,10 @@ defmodule AgentCoordinator.MCPServer do
      "description" => "List all registered codebases",
      "inputSchema" => %{
        "type" => "object",
-        "properties" => %{}
+        "properties" => %{
+          "agent_id" => %{"type" => "string"}
+        },
+        "required" => ["agent_id"]
      }
    },
    %{
@@ -167,12 +176,13 @@ defmodule AgentCoordinator.MCPServer do
      "inputSchema" => %{
        "type" => "object",
        "properties" => %{
+          "agent_id" => %{"type" => "string"},
          "source_codebase_id" => %{"type" => "string"},
          "target_codebase_id" => %{"type" => "string"},
          "dependency_type" => %{"type" => "string"},
          "metadata" => %{"type" => "object"}
        },
-        "required" => ["source_codebase_id", "target_codebase_id", "dependency_type"]
+        "required" => ["agent_id", "source_codebase_id", "target_codebase_id", "dependency_type"]
      }
    },
    %{
@@ -282,6 +292,7 @@ defmodule AgentCoordinator.MCPServer do
      "inputSchema" => %{
        "type" => "object",
        "properties" => %{
+          "agent_id" => %{"type" => "string"},
          "codebase_id" => %{
            "type" => "string",
            "description" => "Optional: filter by codebase ID"
@@ -291,7 +302,8 @@ defmodule AgentCoordinator.MCPServer do
            "default" => true,
            "description" => "Include full task details"
          }
-        }
+        },
+        "required" => ["agent_id"]
      }
    },
    %{
@@ -383,14 +395,23 @@ defmodule AgentCoordinator.MCPServer do
    # Extract agent context for automatic heartbeat management
    case extract_agent_context(request, from, state) do
      {:error, error_message} ->
-        # Return error if agent context extraction fails (unless this is register_agent)
+        # Return error if agent context extraction fails (unless this is an allowed system call)
        method = Map.get(request, "method")
-        if method == "tools/call" and 
-           Map.get(request, "params", %{}) |> Map.get("name") == "register_agent" do
-          # Allow register_agent to proceed without agent_id
+        tool_name = Map.get(request, "params", %{}) |> Map.get("name")
+
+        # Allow certain MCP system calls and register_agent to proceed without agent_id
+        allowed_without_agent = method in ["initialize", "tools/list", "notifications/initialized"] or
+          (method == "tools/call" and tool_name == "register_agent")
+
+        IO.puts(:stderr, "#{method} #{inspect(request)} #{tool_name}")
+
+        if allowed_without_agent do
+          # Allow system calls and register_agent to proceed without agent_id
          response = process_mcp_request(request)
          {:reply, response, state}
        else
+          # Log the rejected call for debugging
+          Logger.warning("Rejected call without agent_id: method=#{method}, tool=#{tool_name}")
          error_response = %{
            "jsonrpc" => "2.0",
            "id" => Map.get(request, "id"),
@@ -402,6 +423,11 @@ defmodule AgentCoordinator.MCPServer do
          {:reply, error_response, state}
        end

+      %{agent_id: nil} ->
+        # System call without agent context
+        response = process_mcp_request(request)
+        {:reply, response, state}
+
      agent_context ->
        # Send pre-operation heartbeat if we have agent context
        if agent_context[:agent_id] do
@@ -478,6 +504,25 @@ defmodule AgentCoordinator.MCPServer do
    }
  end

+  defp process_mcp_request(%{"method" => "notifications/initialized"} = request) do
+    # Handle the initialized notification - this is sent by clients after initialization
+    id = Map.get(request, "id", nil)
+
+    Logger.info("Client initialization notification received")
+
+    # For notifications, we typically don't send a response, but if there's an ID, respond
+    if id do
+      %{
+        "jsonrpc" => "2.0",
+        "id" => id,
+        "result" => %{"status" => "acknowledged"}
+      }
+    else
+      # For notifications without ID, no response is expected
+      nil
+    end
+  end
+
  defp process_mcp_request(
         %{
           "method" => "tools/call",
@@ -540,7 +585,7 @@ defmodule AgentCoordinator.MCPServer do
        case Inbox.start_link(agent.id) do
          {:ok, _pid} -> :ok
          {:error, {:already_started, _pid}} -> :ok
-          {:error, reason} -> 
+          {:error, reason} ->
            Logger.warning("Failed to start inbox for agent #{agent.id}: #{inspect(reason)}")
            :ok
        end
@@ -1382,18 +1427,25 @@ defmodule AgentCoordinator.MCPServer do
  end

  defp extract_agent_context(request, _from, _state) do
-    # Try to get agent_id from various sources
-    cond do
-      # From request arguments
-      Map.get(request, "params", %{})
-      |> Map.get("arguments", %{})
-      |> Map.get("agent_id") ->
-        agent_id = request["params"]["arguments"]["agent_id"]
-        %{agent_id: agent_id}
+    method = Map.get(request, "method")

-      # If no explicit agent_id, return error - agents must register first
-      true ->
-        {:error, "Missing agent_id. Agents must register themselves using register_agent before calling other tools."}
+    # For system calls, don't require agent_id
+    if method in ["initialize", "tools/list", "notifications/initialized"] do
+      %{agent_id: nil}  # System call, no agent context needed
+    else
+      # Try to get agent_id from various sources for non-system calls
+      cond do
+        # From request arguments
+        Map.get(request, "params", %{})
+        |> Map.get("arguments", %{})
+        |> Map.get("agent_id") ->
+          agent_id = request["params"]["arguments"]["agent_id"]
+          %{agent_id: agent_id}
+
+        # If no explicit agent_id, return error - agents must register first
+        true ->
+          {:error, "Missing agent_id. Agents must register themselves using register_agent before calling other tools."}
+      end
    end
  end

--- a/lib/agent_coordinator/task_registry.ex
+++ b/lib/agent_coordinator/task_registry.ex
@@ -340,9 +340,6 @@ defmodule AgentCoordinator.TaskRegistry do
                # Remove from pending since it was assigned
                final_state = %{final_state | pending_tasks: state.pending_tasks}
                {:reply, {:ok, task}, final_state}
-
-              error ->
-                error
            end

          _conflicts ->
@@ -688,8 +685,6 @@ defmodule AgentCoordinator.TaskRegistry do
      :ok -> :ok
      # Inbox already stopped
      {:error, :not_found} -> :ok
-      # Continue regardless
-      _ -> :ok
    end

    # Publish unregistration event
--- a/test_vscode_init.exs
+++ b/test_vscode_init.exs
@@ -0,0 +1,79 @@
+#!/usr/bin/env elixir
+
+# Test script to simulate VS Code MCP initialization sequence
+
+# Start the application
+Application.start(:agent_coordinator)
+
+# Wait a moment for the server to fully start
+Process.sleep(1000)
+
+# Test 1: Initialize call (system call, should work without agent_id)
+IO.puts("Testing initialize call...")
+init_request = %{
+  "jsonrpc" => "2.0",
+  "id" => 1,
+  "method" => "initialize",
+  "params" => %{
+    "protocolVersion" => "2024-11-05",
+    "capabilities" => %{
+      "tools" => %{}
+    },
+    "clientInfo" => %{
+      "name" => "vscode",
+      "version" => "1.0.0"
+    }
+  }
+}
+
+init_response = GenServer.call(AgentCoordinator.MCPServer, {:mcp_request, init_request})
+IO.puts("Initialize response: #{inspect(init_response)}")
+
+# Test 2: Tools/list call (system call, should work without agent_id)
+IO.puts("\nTesting tools/list call...")
+tools_request = %{
+  "jsonrpc" => "2.0",
+  "id" => 2,
+  "method" => "tools/list"
+}
+
+tools_response = GenServer.call(AgentCoordinator.MCPServer, {:mcp_request, tools_request})
+IO.puts("Tools/list response: #{inspect(tools_response)}")
+
+# Test 3: Register agent call (should work)
+IO.puts("\nTesting register_agent call...")
+register_request = %{
+  "jsonrpc" => "2.0",
+  "id" => 3,
+  "method" => "tools/call",
+  "params" => %{
+    "name" => "register_agent",
+    "arguments" => %{
+      "name" => "GitHub Copilot Test Agent",
+      "capabilities" => ["file_operations", "code_generation"]
+    }
+  }
+}
+
+register_response = GenServer.call(AgentCoordinator.MCPServer, {:mcp_request, register_request})
+IO.puts("Register agent response: #{inspect(register_response)}")
+
+# Test 4: Try a call that requires agent_id (should fail without agent_id)
+IO.puts("\nTesting call that requires agent_id (should fail)...")
+task_request = %{
+  "jsonrpc" => "2.0",
+  "id" => 4,
+  "method" => "tools/call",
+  "params" => %{
+    "name" => "create_task",
+    "arguments" => %{
+      "title" => "Test task",
+      "description" => "This should fail without agent_id"
+    }
+  }
+}
+
+task_response = GenServer.call(AgentCoordinator.MCPServer, {:mcp_request, task_request})
+IO.puts("Task creation response: #{inspect(task_response)}")
+
+IO.puts("\n✅ All tests completed!")"