mulit-agent/controller/internal/controlplane/mcp_test.go

package controlplane

import (
	"path/filepath"
	"testing"
)

func TestMCPToolRepositoryValidatesDefinitions(t *testing.T) {
	t.Setenv("MULTICLAW_CONTROLLER_DATA_DIR", t.TempDir())
	controller, err := Open(filepath.Join(t.TempDir(), "controller-config.json"))
	if err != nil {
		t.Fatal(err)
	}
	defer controller.Close()

	// The legacy narrow tool is seeded as the only builtin an Agent implements.
	tools, err := controller.MCPTools()
	if err != nil || len(tools) != 1 || tools[0].ID != "desktop-filesystem" || tools[0].Kind != toolKindBuiltin {
		t.Fatalf("expected the seeded builtin tool, got %#v / %v", tools, err)
	}
	if _, err := controller.SaveMCPTool(MCPToolInput{ID: "docker", Name: "Docker", Kind: toolKindBuiltin}); err == nil {
		t.Fatal("an unknown builtin tool must be rejected")
	}
	if _, err := controller.SaveMCPTool(MCPToolInput{ID: "imageforge", Name: "ImageForge", Kind: toolKindHTTP, Endpoint: "not-a-url"}); err == nil {
		t.Fatal("an invalid MCP endpoint must be rejected")
	}
	if _, err := controller.SaveMCPTool(MCPToolInput{ID: "Bad ID", Name: "x", Kind: toolKindHTTP, Endpoint: "http://127.0.0.1:9099/mcp"}); err == nil {
		t.Fatal("an invalid tool identifier must be rejected")
	}
	tool, err := controller.SaveMCPTool(MCPToolInput{ID: "imageforge", Name: "ImageForge MCP", Description: "本地桌面控制", Kind: toolKindHTTP, Endpoint: "http://127.0.0.1:9099/mcp/"})
	if err != nil {
		t.Fatal(err)
	}
	if tool.Endpoint != "http://127.0.0.1:9099/mcp" || tool.Version != "1.0.0" || tool.AgentCount != 0 {
		t.Fatalf("unexpected saved tool: %#v", tool)
	}
}

func TestMCPToolRepositoryValidatesStdioCommands(t *testing.T) {
	t.Setenv("MULTICLAW_CONTROLLER_DATA_DIR", t.TempDir())
	controller, err := Open(filepath.Join(t.TempDir(), "controller-config.json"))
	if err != nil {
		t.Fatal(err)
	}
	defer controller.Close()

	if _, err := controller.SaveMCPTool(MCPToolInput{ID: "local-plugin", Name: "本地插件", Kind: toolKindStdio}); err == nil {
		t.Fatal("a stdio tool without a command must be rejected")
	}
	if _, err := controller.SaveMCPTool(MCPToolInput{ID: "local-plugin", Name: "本地插件", Kind: toolKindStdio, Command: "python", Args: []string{"line\nbreak"}}); err == nil {
		t.Fatal("an argument containing a newline must be rejected")
	}
	if _, err := controller.SaveMCPTool(MCPToolInput{ID: "local-plugin", Name: "本地插件", Kind: toolKindStdio, Command: "python", Args: make([]string, maxToolArguments+1)}); err == nil {
		t.Fatal("too many arguments must be rejected")
	}

	tool, err := controller.SaveMCPTool(MCPToolInput{
		ID: "local-plugin", Name: "本地插件", Kind: toolKindStdio,
		Command: "python", Args: []string{"-m", "my_mcp_server", "--stdio"}, Version: "1.0.0",
	})
	if err != nil {
		t.Fatal(err)
	}
	if tool.Command != "python" || len(tool.Args) != 3 || tool.Endpoint != "" {
		t.Fatalf("unexpected stdio tool: %#v", tool)
	}
	// The command line must survive a reload, since it is what the Agent spawns.
	stored, err := controller.MCPTools()
	if err != nil {
		t.Fatal(err)
	}
	persisted := findTool(t, stored, "local-plugin")
	if persisted.Command != "python" || len(persisted.Args) != 3 || persisted.Args[2] != "--stdio" {
		t.Fatalf("expected the persisted stdio command line, got %#v", persisted)
	}
	// A rename with an unchanged command line must not create a second tool.
	if _, err := controller.SaveMCPTool(MCPToolInput{ID: "local-plugin", Name: "本地插件（改名）", Kind: toolKindStdio, Command: "python", Args: []string{"-m", "my_mcp_server", "--stdio"}, Version: "1.0.0"}); err != nil {
		t.Fatal(err)
	}
	tools, err := controller.MCPTools()
	if err != nil || len(tools) != 2 || findTool(t, tools, "local-plugin").Name != "本地插件（改名）" {
		t.Fatalf("unexpected repository state: %#v / %v", tools, err)
	}
}

func findTool(t *testing.T, tools []MCPTool, toolID string) MCPTool {
	t.Helper()
	for _, tool := range tools {
		if tool.ID == toolID {
			return tool
		}
	}
	t.Fatalf("tool %q not found in %#v", toolID, tools)
	return MCPTool{}
}

func TestMCPToolAssignmentBumpsOnlyChangedDefinitions(t *testing.T) {
	t.Setenv("MULTICLAW_CONTROLLER_DATA_DIR", t.TempDir())
	controller, err := Open(filepath.Join(t.TempDir(), "controller-config.json"))
	if err != nil {
		t.Fatal(err)
	}
	defer controller.Close()
	if _, err := controller.db.Exec("INSERT INTO managed_agents(agent_id, name, credential_hash, status, last_seen_at) VALUES ('agent-1', '测试机', 'hash', 'offline', '2026-01-01T00:00:00Z')"); err != nil {
		t.Fatal(err)
	}
	if _, err := controller.SaveMCPTool(MCPToolInput{ID: "imageforge", Name: "ImageForge MCP", Kind: toolKindHTTP, Endpoint: "http://127.0.0.1:9099/mcp"}); err != nil {
		t.Fatal(err)
	}
	if err := controller.SetAgentMCPTools("agent-1", []string{"imageforge"}); err != nil {
		t.Fatal(err)
	}
	if version := assignmentVersion(t, controller, "agent-1", "imageforge"); version != 1 {
		t.Fatalf("expected the first deployment version, got %d", version)
	}
	if err := controller.SetAgentMCPTools("agent-1", []string{"imageforge"}); err != nil {
		t.Fatal(err)
	}
	if version := assignmentVersion(t, controller, "agent-1", "imageforge"); version != 1 {
		t.Fatalf("re-selecting an unchanged tool must not bump the version, got %d", version)
	}

	// A changed definition must be re-applied instead of being rejected as stale.
	if _, err := controller.SaveMCPTool(MCPToolInput{ID: "imageforge", Name: "ImageForge MCP", Kind: toolKindHTTP, Endpoint: "http://127.0.0.1:9100/mcp", Version: "1.1.0"}); err != nil {
		t.Fatal(err)
	}
	if version := assignmentVersion(t, controller, "agent-1", "imageforge"); version != 2 {
		t.Fatalf("expected the changed definition to be redeployed, got version %d", version)
	}
	statuses, err := controller.AgentMCPStatuses()
	if err != nil || len(statuses) != 1 || statuses[0].ToolID != "imageforge" || statuses[0].ApplyStatus != "pending" || statuses[0].Endpoint != "http://127.0.0.1:9100/mcp" {
		t.Fatalf("unexpected tool deployment status: %#v / %v", statuses, err)
	}
}

func TestDeleteMCPToolWithdrawsEveryAssignment(t *testing.T) {
	t.Setenv("MULTICLAW_CONTROLLER_DATA_DIR", t.TempDir())
	controller, err := Open(filepath.Join(t.TempDir(), "controller-config.json"))
	if err != nil {
		t.Fatal(err)
	}
	defer controller.Close()
	if _, err := controller.db.Exec("INSERT INTO managed_agents(agent_id, name, credential_hash, status, last_seen_at) VALUES ('agent-1', '测试机', 'hash', 'offline', '2026-01-01T00:00:00Z')"); err != nil {
		t.Fatal(err)
	}
	if err := controller.SetAgentMCPTools("agent-1", []string{"desktop-filesystem"}); err != nil {
		t.Fatal(err)
	}
	if err := controller.DeleteMCPTool("desktop-filesystem"); err != nil {
		t.Fatal(err)
	}
	if err := controller.DeleteMCPTool("desktop-filesystem"); err == nil {
		t.Fatal("deleting a missing tool must fail")
	}
	var assignments int
	if err := controller.db.QueryRow("SELECT COUNT(*) FROM agent_mcp_assignments").Scan(&assignments); err != nil || assignments != 0 {
		t.Fatalf("expected every assignment to be withdrawn, got %d / %v", assignments, err)
	}
	if tools, err := controller.MCPTools(); err != nil || len(tools) != 0 {
		t.Fatalf("expected an empty tool repository, got %#v / %v", tools, err)
	}
}

func assignmentVersion(t *testing.T, controller *ControlPlane, agentID, toolID string) int {
	t.Helper()
	var version int
	if err := controller.db.QueryRow("SELECT deployment_version FROM agent_mcp_assignments WHERE agent_id = ? AND tool_id = ?", agentID, toolID).Scan(&version); err != nil {
		t.Fatal(err)
	}
	return version
}