Compare commits

...
1 Commits
Author SHA1 Message Date
ugurtafrali 1ea8bad074 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.
2026-04-17 14:40:41 +08:00
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},