From 56af1ea091165113a3edb73672b5d82ec8affc8b Mon Sep 17 00:00:00 2001 From: ybc Date: Tue, 14 Apr 2026 19:24:47 +0800 Subject: [PATCH] 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 --- internal/app/auth_registry_test.go | 213 +++++++++++++++++++++++++++++ internal/app/root.go | 52 +++++-- internal/plugin/plugin_test.go | 122 +++++++++++++++++ 3 files changed, 377 insertions(+), 10 deletions(-) create mode 100644 internal/app/auth_registry_test.go diff --git a/internal/app/auth_registry_test.go b/internal/app/auth_registry_test.go new file mode 100644 index 00000000..2fbfac92 --- /dev/null +++ b/internal/app/auth_registry_test.go @@ -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") + } +} diff --git a/internal/app/root.go b/internal/app/root.go index 0a7aef68..d92a01a9 100644 --- a/internal/app/root.go +++ b/internal/app/root.go @@ -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...) } } diff --git a/internal/plugin/plugin_test.go b/internal/plugin/plugin_test.go index 06e44238..2e38846a 100644 --- a/internal/plugin/plugin_test.go +++ b/internal/plugin/plugin_test.go @@ -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{