From 1ea8bad074a83fc88baffeefbc7d9a02a7b886ca Mon Sep 17 00:00:00 2001 From: ugurtafrali Date: Fri, 17 Apr 2026 05:12:53 +0300 Subject: [PATCH] Reduce MCP discovery timeouts When MCP plugins are unreachable, the discovery times out too slowly and blocks commands. Reduced the overall discovery timeout from 10s to 4s and per-server timeout from 5s to 2s. Made the per-server timeout configurable. --- internal/cli/loader.go | 2 +- internal/discovery/service.go | 21 +++++++----- internal/discovery/service_test.go | 55 ++++++++++++++++++++++++++++++ internal/transport/client.go | 2 +- 4 files changed, 70 insertions(+), 10 deletions(-) diff --git a/internal/cli/loader.go b/internal/cli/loader.go index d90c4a25..4527f25b 100644 --- a/internal/cli/loader.go +++ b/internal/cli/loader.go @@ -107,7 +107,7 @@ const ( DefaultMarketBaseURL = "https://mcp.dingtalk.com" // defaultDiscoveryTimeout bounds the time spent on live registry discovery. - defaultDiscoveryTimeout = 10 * time.Second + defaultDiscoveryTimeout = 4 * time.Second ) type CatalogLoader interface { diff --git a/internal/discovery/service.go b/internal/discovery/service.go index c30d2713..7cb570c4 100644 --- a/internal/discovery/service.go +++ b/internal/discovery/service.go @@ -53,12 +53,13 @@ const ( var errCLIServerSkipped = errors.New("server marked cli.skip") type Service struct { - MarketClient *market.Client - Transport *transport.Client - Cache *cache.Store - Tenant string - AuthIdentity string - Logger *slog.Logger + MarketClient *market.Client + Transport *transport.Client + Cache *cache.Store + Tenant string + AuthIdentity string + Logger *slog.Logger + PerServerTimeout time.Duration // overrides perServerDiscoveryTimeout when > 0 } type RuntimeServer struct { @@ -170,7 +171,7 @@ func (s *Service) DiscoverServerRuntime(ctx context.Context, server market.Serve }, nil } -const perServerDiscoveryTimeout = 5 * time.Second +const perServerDiscoveryTimeout = 2 * time.Second func (s *Service) DiscoverAllRuntime(ctx context.Context, servers []market.ServerDescriptor) ([]RuntimeServer, []RuntimeFailure) { type discoveryResult struct { @@ -189,12 +190,16 @@ func (s *Service) DiscoverAllRuntime(ctx context.Context, servers []market.Serve } ch := make(chan discoveryResult, len(filtered)) + serverTimeout := s.PerServerTimeout + if serverTimeout <= 0 { + serverTimeout = perServerDiscoveryTimeout + } var wg sync.WaitGroup for _, srv := range filtered { wg.Add(1) go func(server market.ServerDescriptor) { defer wg.Done() - serverCtx, cancel := context.WithTimeout(ctx, perServerDiscoveryTimeout) + serverCtx, cancel := context.WithTimeout(ctx, serverTimeout) defer cancel() start := time.Now() rs, err := s.DiscoverServerRuntime(serverCtx, server) diff --git a/internal/discovery/service_test.go b/internal/discovery/service_test.go index c4d729a7..00f40b87 100644 --- a/internal/discovery/service_test.go +++ b/internal/discovery/service_test.go @@ -6,6 +6,7 @@ import ( "net/http" "net/http/httptest" "testing" + "time" "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/cache" "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/market" @@ -435,6 +436,60 @@ func TestParseDetailSchema(t *testing.T) { } } +func TestDiscoverAllRuntime_TimeoutFallsBackToCache(t *testing.T) { + t.Parallel() + + // done signals slow handlers to exit so srv.Close() can complete. + done := make(chan struct{}) + // Server that blocks until signalled (simulates an unreachable MCP server). + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-done: + case <-time.After(10 * time.Second): + } + http.Error(w, "timeout", http.StatusServiceUnavailable) + })) + // LIFO: close(done) runs first so handlers exit, then srv.Close() completes. + defer srv.Close() + defer close(done) + + svc := newTestService(t, srv.URL, srv) + // Override timeout so the test completes quickly. + svc.PerServerTimeout = 80 * time.Millisecond + + server := market.ServerDescriptor{ + Key: "slow-server", + Endpoint: srv.URL + "/mcp", + } + + // Pre-populate the cache so the fallback has data to return. + partition := "test-tenant/test-identity" + _ = svc.Cache.SaveTools(partition, server.Key, cache.ToolsSnapshot{ + ServerKey: server.Key, + ProtocolVersion: "2025-03-26", + Tools: []transport.ToolDescriptor{ + {Name: "cached-tool", Description: "from cache"}, + }, + }) + + results, failures := svc.DiscoverAllRuntime(context.Background(), []market.ServerDescriptor{server}) + if len(failures) != 0 { + t.Fatalf("failures count = %d, want 0 (expected cache fallback)", len(failures)) + } + if len(results) != 1 { + t.Fatalf("results count = %d, want 1", len(results)) + } + if !results[0].Degraded { + t.Fatal("cache fallback result should be degraded") + } + if len(results[0].Tools) == 0 { + t.Fatal("expected cached tools to be returned") + } + if results[0].Tools[0].Name != "cached-tool" { + t.Fatalf("tool name = %q, want cached-tool", results[0].Tools[0].Name) + } +} + func TestPartition(t *testing.T) { t.Parallel() svc := &Service{Tenant: "corp1", AuthIdentity: "user1"} diff --git a/internal/transport/client.go b/internal/transport/client.go index b5ae8c74..5154c82c 100644 --- a/internal/transport/client.go +++ b/internal/transport/client.go @@ -212,7 +212,7 @@ func (r *ToolCallResult) UnmarshalJSON(data []byte) error { func defaultTransport() *http.Transport { return &http.Transport{ DialContext: (&net.Dialer{ - Timeout: 10 * time.Second, + Timeout: 3 * time.Second, KeepAlive: 30 * time.Second, }).DialContext, TLSClientConfig: &tls.Config{MinVersion: tls.VersionTLS12},