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.
This commit is contained in:
ugurtafrali
2026-04-17 14:40:41 +08:00
committed by 修雨
parent 7b7aeadbbe
commit 1ea8bad074
4 changed files with 70 additions and 10 deletions
+1 -1
View File
@@ -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 {
+13 -8
View File
@@ -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)
+55
View File
@@ -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"}
+1 -1
View File
@@ -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},