feat: parallel service discovery + unit tests for auth-headers

- Optimize loadPlugins to discover HTTP servers in parallel (sync.WaitGroup) when multiple remote servers exist, reducing startup time from sequential N*10s to parallel max(10s) - Add auth_registry_test.go with 7 unit tests covering: - PluginAuth registry CRUD and multi-product isolation - deriveToolCLIName conversion (6 cases) - registerPluginAuthFromHeaders extraction and registration - buildPluginAuthClient with/without Authorization header - Add 3 test cases to plugin_test.go covering: - Headers -> AuthHeaders conversion with env var expansion - No-headers descriptor produces nil AuthHeaders - ParseManifest correctly parses headers field from plugin.json
2026-04-16 09:38:18 +08:00
parent ea5859b92b
commit 56af1ea091
3 changed files with 377 additions and 10 deletions
@@ -0,0 +1,213 @@
+// Copyright 2026 Alibaba Group
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package app
+
+import (
+	"testing"
+
+	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/market"
+	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/transport"
+)
+
+func TestPluginAuthRegistry(t *testing.T) {
+	// Clean up after test
+	defer func() {
+		pluginAuthMu.Lock()
+		delete(pluginAuthRegistry, "test-product")
+		pluginAuthMu.Unlock()
+	}()
+
+	// Initially not found
+	if _, ok := LookupPluginAuth("test-product"); ok {
+		t.Error("expected LookupPluginAuth to return false for unregistered product")
+	}
+
+	// Register auth credentials
+	auth := &PluginAuth{
+		Token:          "sk-test-token-12345",
+		ExtraHeaders:   map[string]string{"X-Custom": "value"},
+		TrustedDomains: []string{"api.example.com", "*.example.com"},
+	}
+	RegisterPluginAuth("test-product", auth)
+
+	// Now should be found
+	got, ok := LookupPluginAuth("test-product")
+	if !ok {
+		t.Fatal("expected LookupPluginAuth to return true after registration")
+	}
+	if got != auth {
+		t.Error("LookupPluginAuth returned different auth instance")
+	}
+	if got.Token != "sk-test-token-12345" {
+		t.Errorf("Token = %q, want sk-test-token-12345", got.Token)
+	}
+	if got.ExtraHeaders["X-Custom"] != "value" {
+		t.Errorf("ExtraHeaders[X-Custom] = %q, want value", got.ExtraHeaders["X-Custom"])
+	}
+	if len(got.TrustedDomains) != 2 {
+		t.Errorf("TrustedDomains len = %d, want 2", len(got.TrustedDomains))
+	}
+}
+
+func TestPluginAuthRegistryIsolation(t *testing.T) {
+	// Clean up after test
+	defer func() {
+		pluginAuthMu.Lock()
+		delete(pluginAuthRegistry, "product-a")
+		delete(pluginAuthRegistry, "product-b")
+		pluginAuthMu.Unlock()
+	}()
+
+	authA := &PluginAuth{Token: "token-a"}
+	authB := &PluginAuth{Token: "token-b"}
+
+	RegisterPluginAuth("product-a", authA)
+	RegisterPluginAuth("product-b", authB)
+
+	gotA, okA := LookupPluginAuth("product-a")
+	gotB, okB := LookupPluginAuth("product-b")
+
+	if !okA || !okB {
+		t.Fatal("expected both products to be registered")
+	}
+	if gotA.Token != "token-a" {
+		t.Errorf("product-a Token = %q, want token-a", gotA.Token)
+	}
+	if gotB.Token != "token-b" {
+		t.Errorf("product-b Token = %q, want token-b", gotB.Token)
+	}
+}
+
+func TestDeriveToolCLIName(t *testing.T) {
+	tests := []struct {
+		input string
+		want  string
+	}{
+		{"web_search", "web-search"},
+		{"maps.search_poi", "search-poi"},
+		{"maps.geo", "geo"},
+		{"simple", "simple"},
+		{"a.b.deep_nested_name", "deep-nested-name"},
+		{"already-kebab", "already-kebab"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.input, func(t *testing.T) {
+			got := deriveToolCLIName(tt.input)
+			if got != tt.want {
+				t.Errorf("deriveToolCLIName(%q) = %q, want %q", tt.input, got, tt.want)
+			}
+		})
+	}
+}
+
+func TestRegisterPluginAuthFromHeaders(t *testing.T) {
+	// Clean up after test
+	defer func() {
+		pluginAuthMu.Lock()
+		delete(pluginAuthRegistry, "test-srv")
+		pluginAuthMu.Unlock()
+	}()
+
+	srv := market.ServerDescriptor{
+		Key:      "test-srv",
+		Endpoint: "https://api.example.com/mcp/v1",
+		CLI:      market.CLIOverlay{ID: "test-srv", Command: "test-srv"},
+		AuthHeaders: map[string]string{
+			"Authorization": "Bearer sk-my-secret-key",
+			"X-Custom":      "custom-value",
+		},
+	}
+
+	registerPluginAuthFromHeaders(srv)
+
+	auth, ok := LookupPluginAuth("test-srv")
+	if !ok {
+		t.Fatal("expected auth to be registered after registerPluginAuthFromHeaders")
+	}
+	if auth.Token != "sk-my-secret-key" {
+		t.Errorf("Token = %q, want sk-my-secret-key", auth.Token)
+	}
+	if auth.ExtraHeaders["X-Custom"] != "custom-value" {
+		t.Errorf("ExtraHeaders[X-Custom] = %q, want custom-value", auth.ExtraHeaders["X-Custom"])
+	}
+	if len(auth.TrustedDomains) != 2 {
+		t.Fatalf("TrustedDomains len = %d, want 2", len(auth.TrustedDomains))
+	}
+	if auth.TrustedDomains[0] != "api.example.com" {
+		t.Errorf("TrustedDomains[0] = %q, want api.example.com", auth.TrustedDomains[0])
+	}
+}
+
+func TestRegisterPluginAuthFromHeadersNoAuth(t *testing.T) {
+	srv := market.ServerDescriptor{
+		Key:      "no-auth-srv",
+		Endpoint: "https://api.example.com/mcp/v1",
+		CLI:      market.CLIOverlay{ID: "no-auth-srv"},
+		AuthHeaders: map[string]string{
+			"X-Custom": "custom-value",
+		},
+	}
+
+	registerPluginAuthFromHeaders(srv)
+
+	// Should not register because there's no Authorization header
+	if _, ok := LookupPluginAuth("no-auth-srv"); ok {
+		t.Error("expected no auth registration when Authorization header is missing")
+	}
+}
+
+func TestBuildPluginAuthClient(t *testing.T) {
+	base := transport.NewClient(nil)
+
+	srv := market.ServerDescriptor{
+		Endpoint: "https://dashscope.aliyuncs.com/compatible-mode/v1/mcp",
+		AuthHeaders: map[string]string{
+			"Authorization": "Bearer sk-test-api-key",
+			"X-Extra":       "extra-value",
+		},
+	}
+
+	client := buildPluginAuthClient(base, srv)
+
+	// Should return a different client instance
+	if client == base {
+		t.Error("expected buildPluginAuthClient to return a new client, not the base")
+	}
+
+	// Verify trusted domains
+	if len(client.TrustedDomains) != 2 {
+		t.Fatalf("TrustedDomains len = %d, want 2", len(client.TrustedDomains))
+	}
+	if client.TrustedDomains[0] != "dashscope.aliyuncs.com" {
+		t.Errorf("TrustedDomains[0] = %q, want dashscope.aliyuncs.com", client.TrustedDomains[0])
+	}
+}
+
+func TestBuildPluginAuthClientNoAuth(t *testing.T) {
+	base := transport.NewClient(nil)
+
+	srv := market.ServerDescriptor{
+		Endpoint: "https://api.example.com/mcp/v1",
+		AuthHeaders: map[string]string{
+			"X-Custom": "custom-value",
+		},
+	}
+
+	client := buildPluginAuthClient(base, srv)
+
+	// Should return the base client when no Authorization header
+	if client != base {
+		t.Error("expected buildPluginAuthClient to return base client when no Authorization header")
+	}
+}
@@ -26,6 +26,7 @@ import (
 	"path"
 	"path/filepath"
 	"strings"
+	"sync"
 	"syscall"
 	"time"

@@ -984,25 +985,56 @@ func loadPlugins(engine *pipeline.Engine, runner executor.Runner) []*cobra.Comma
 	allPlugins := append(managedPlugins, userPlugins...)
 	allPlugins = append(allPlugins, devPlugins...)

-	// 3. Discover tools from streamable-http servers and build CLI commands
+	// 3. Discover tools from streamable-http servers and build CLI commands.
+	//    Third-party servers with auth headers are discovered in parallel
+	//    to avoid sequential 10s timeouts when multiple remote servers exist.
 	var pluginCmds []*cobra.Command
 	tc := transport.NewClient(nil)
+
+	// Collect all server descriptors and register auth first (fast, no I/O).
+	type pluginServer struct {
+		plugin *plugin.Plugin
+		srv    market.ServerDescriptor
+	}
+	var httpServers []pluginServer
+
 	for _, p := range allPlugins {
 		for _, srv := range p.ToServerDescriptors() {
 			AppendDynamicServer(srv)

-			// For plugin-owned streamable-http servers, also discover tools
-			// and build CLI commands (Market servers rely on the detail API).
-			if srv.HasCLIMeta {
-				cmds := registerHTTPServer(p, srv, tc, runner)
-				pluginCmds = append(pluginCmds, cmds...)
-			}
-
-			// Register plugin-level auth credentials for third-party servers
-			// so the runner can inject the correct token at execution time.
 			if len(srv.AuthHeaders) > 0 {
 				registerPluginAuthFromHeaders(srv)
 			}
+
+			if srv.HasCLIMeta {
+				httpServers = append(httpServers, pluginServer{plugin: p, srv: srv})
+			}
+		}
+	}
+
+	// Discover tools from HTTP servers in parallel when there are multiple
+	// servers with auth headers (third-party services with higher latency).
+	if len(httpServers) > 1 {
+		type discoveryResult struct {
+			commands []*cobra.Command
+		}
+		results := make([]discoveryResult, len(httpServers))
+		var wg sync.WaitGroup
+		for i, ps := range httpServers {
+			wg.Add(1)
+			go func(idx int, ps pluginServer) {
+				defer wg.Done()
+				results[idx].commands = registerHTTPServer(ps.plugin, ps.srv, tc, runner)
+			}(i, ps)
+		}
+		wg.Wait()
+		for _, r := range results {
+			pluginCmds = append(pluginCmds, r.commands...)
+		}
+	} else {
+		for _, ps := range httpServers {
+			cmds := registerHTTPServer(ps.plugin, ps.srv, tc, runner)
+			pluginCmds = append(pluginCmds, cmds...)
 		}
 	}

@@ -226,6 +226,128 @@ func TestPluginToServerDescriptors(t *testing.T) {
 	}
 }

+func TestPluginToServerDescriptorsWithHeaders(t *testing.T) {
+	cliOverlay, _ := json.Marshal(map[string]any{
+		"id":      "web-search",
+		"command": "web-search",
+	})
+
+	// Set an environment variable to test expansion
+	t.Setenv("TEST_API_KEY", "sk-test-12345")
+
+	p := &Plugin{
+		Manifest: Manifest{
+			Name:        "my-plugin",
+			Description: "Test plugin with headers",
+			MCPServers: map[string]*MCPServer{
+				"web-search": {
+					Type:     "streamable-http",
+					Endpoint: "https://api.example.com/mcp/v1",
+					CLI:      cliOverlay,
+					Headers: map[string]string{
+						"Authorization": "Bearer ${TEST_API_KEY}",
+						"X-Custom":      "static-value",
+					},
+				},
+			},
+		},
+		Root: "/tmp/plugins/my-plugin",
+	}
+
+	descriptors := p.ToServerDescriptors()
+	if len(descriptors) != 1 {
+		t.Fatalf("got %d descriptors, want 1", len(descriptors))
+	}
+
+	d := descriptors[0]
+	if d.Key != "web-search" {
+		t.Errorf("key = %q, want web-search", d.Key)
+	}
+	if len(d.AuthHeaders) != 2 {
+		t.Fatalf("AuthHeaders len = %d, want 2", len(d.AuthHeaders))
+	}
+	// Environment variable should be expanded
+	if d.AuthHeaders["Authorization"] != "Bearer sk-test-12345" {
+		t.Errorf("AuthHeaders[Authorization] = %q, want 'Bearer sk-test-12345'", d.AuthHeaders["Authorization"])
+	}
+	if d.AuthHeaders["X-Custom"] != "static-value" {
+		t.Errorf("AuthHeaders[X-Custom] = %q, want static-value", d.AuthHeaders["X-Custom"])
+	}
+	if d.Source != "plugin" {
+		t.Errorf("source = %q, want plugin (non-managed)", d.Source)
+	}
+}
+
+func TestPluginToServerDescriptorsNoHeaders(t *testing.T) {
+	cliOverlay, _ := json.Marshal(map[string]any{
+		"id":      "conference",
+		"command": "conference",
+	})
+
+	p := &Plugin{
+		Manifest: Manifest{
+			Name: "conference",
+			MCPServers: map[string]*MCPServer{
+				"conference": {
+					Type:     "streamable-http",
+					Endpoint: "https://mcp.conference.dingtalk.com",
+					CLI:      cliOverlay,
+				},
+			},
+		},
+		Root: "/tmp/plugins/conference",
+	}
+
+	descriptors := p.ToServerDescriptors()
+	if len(descriptors) != 1 {
+		t.Fatalf("got %d descriptors, want 1", len(descriptors))
+	}
+	if descriptors[0].AuthHeaders != nil {
+		t.Errorf("AuthHeaders = %v, want nil for server without headers", descriptors[0].AuthHeaders)
+	}
+}
+
+func TestParseManifestWithHeaders(t *testing.T) {
+	dir := t.TempDir()
+	manifestPath := filepath.Join(dir, "plugin.json")
+	content := `{
+		"name": "api-plugin",
+		"version": "1.0.0",
+		"mcpServers": {
+			"api-server": {
+				"type": "streamable-http",
+				"endpoint": "https://api.example.com/mcp",
+				"headers": {
+					"Authorization": "Bearer ${MY_API_KEY}",
+					"X-Custom-Header": "custom-value"
+				}
+			}
+		}
+	}`
+	if err := os.WriteFile(manifestPath, []byte(content), 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	m, err := ParseManifest(manifestPath)
+	if err != nil {
+		t.Fatalf("ParseManifest: %v", err)
+	}
+
+	srv := m.MCPServers["api-server"]
+	if srv == nil {
+		t.Fatal("api-server not found in MCPServers")
+	}
+	if len(srv.Headers) != 2 {
+		t.Fatalf("Headers len = %d, want 2", len(srv.Headers))
+	}
+	if srv.Headers["Authorization"] != "Bearer ${MY_API_KEY}" {
+		t.Errorf("Headers[Authorization] = %q, want raw template", srv.Headers["Authorization"])
+	}
+	if srv.Headers["X-Custom-Header"] != "custom-value" {
+		t.Errorf("Headers[X-Custom-Header] = %q, want custom-value", srv.Headers["X-Custom-Header"])
+	}
+}
+
 func TestLoaderScanEmpty(t *testing.T) {
 	dir := t.TempDir()
 	loader := &Loader{