Skip to content

Commit 22238e4

Browse files
DumbrisDumbris
committed
Enhance timeout management and logging in client and server components
- Updated timeout duration for tool listing operations to 30 seconds for better handling of Docker container startup and API calls. - Improved logging for various operations, including connection errors and tool calls, to provide more detailed insights during debugging. - Added conditional logging for JSON request/response data based on debug level, enhancing traceability without cluttering logs.
1 parent 9272d9a commit 22238e4

3 files changed

Lines changed: 104 additions & 32 deletions

File tree

‎internal/server/server.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -774,7 +774,8 @@ func (s *Server) getServerToolCount(serverID string) int {
774774
return 0
775775
}
776776

777-
ctx := context.Background()
777+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
778+
defer cancel()
778779
tools, err := client.ListTools(ctx)
779780
if err != nil {
780781
s.logger.Warn("Failed to get tool count for server",

‎internal/upstream/client.go‎

Lines changed: 100 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ func NewClient(id string, serverConfig *config.ServerConfig, logger *zap.Logger,
7676
} else {
7777
envConfig = secureenv.DefaultEnvConfig()
7878
}
79-
79+
8080
// Add server-specific environment variables to the custom vars
8181
if len(serverConfig.Env) > 0 {
8282
// Create a copy of the environment config with server-specific variables
@@ -91,15 +91,15 @@ func NewClient(id string, serverConfig *config.ServerConfig, logger *zap.Logger,
9191
}
9292
serverEnvConfig.CustomVars = customVars
9393
}
94-
94+
9595
// Add server-specific environment variables
9696
for k, v := range serverConfig.Env {
9797
serverEnvConfig.CustomVars[k] = v
9898
}
99-
99+
100100
envConfig = &serverEnvConfig
101101
}
102-
102+
103103
c.envManager = secureenv.NewManager(envConfig)
104104

105105
// Create upstream server logger if logging config is provided
@@ -220,6 +220,17 @@ func (c *Client) Connect(ctx context.Context) error {
220220
zap.Int("filtered_count", len(envVars)),
221221
zap.String("command", command))
222222

223+
// Debug logging for stdio transport (only at debug level)
224+
c.logger.Debug("Creating stdio transport",
225+
zap.String("command", command),
226+
zap.Strings("args", cmdArgs),
227+
zap.Int("env_count", len(envVars)))
228+
229+
if c.upstreamLogger != nil {
230+
c.upstreamLogger.Debug("Process starting",
231+
zap.String("full_command", fmt.Sprintf("%s %s", command, strings.Join(cmdArgs, " "))))
232+
}
233+
223234
stdioTransport := transport.NewStdio(command, envVars, cmdArgs...)
224235
c.client = client.NewClient(stdioTransport)
225236
default:
@@ -243,6 +254,21 @@ func (c *Client) Connect(ctx context.Context) error {
243254
c.retryCount++
244255
c.lastRetryTime = time.Now()
245256
c.mu.Unlock()
257+
258+
// Log to both main and server logs for critical errors
259+
c.logger.Error("Failed to start MCP client", zap.Error(err))
260+
if c.upstreamLogger != nil {
261+
c.upstreamLogger.Error("Client start failed", zap.Error(err))
262+
}
263+
264+
// Add debug transport info if DEBUG level is enabled
265+
if c.logger.Core().Enabled(zap.DebugLevel) {
266+
c.logger.Debug("MCP client start failed details",
267+
zap.String("error_type", fmt.Sprintf("%T", err)),
268+
zap.String("command", c.config.Command),
269+
zap.Strings("args", c.config.Args))
270+
}
271+
246272
return fmt.Errorf("failed to start MCP client: %w", err)
247273
}
248274

@@ -262,6 +288,13 @@ func (c *Client) Connect(ctx context.Context) error {
262288
c.retryCount++
263289
c.lastRetryTime = time.Now()
264290
c.mu.Unlock()
291+
292+
// Log to both main and server logs for critical errors
293+
c.logger.Error("Failed to initialize MCP client", zap.Error(err))
294+
if c.upstreamLogger != nil {
295+
c.upstreamLogger.Error("Initialize failed", zap.Error(err))
296+
}
297+
265298
c.client.Close()
266299
return fmt.Errorf("failed to initialize MCP client: %w", err)
267300
}
@@ -277,18 +310,31 @@ func (c *Client) Connect(ctx context.Context) error {
277310
zap.String("server_name", serverInfo.ServerInfo.Name),
278311
zap.String("server_version", serverInfo.ServerInfo.Version))
279312

313+
// Add debug transport info if DEBUG level is enabled
314+
if c.logger.Core().Enabled(zap.DebugLevel) {
315+
c.logger.Debug("MCP connection details",
316+
zap.String("protocol_version", serverInfo.ProtocolVersion),
317+
zap.String("command", c.config.Command),
318+
zap.Strings("args", c.config.Args),
319+
zap.String("transport", c.determineTransportType()))
320+
}
321+
280322
if c.upstreamLogger != nil {
281323
c.upstreamLogger.Info("Connected successfully",
282324
zap.String("server_name", serverInfo.ServerInfo.Name),
283325
zap.String("server_version", serverInfo.ServerInfo.Version),
284326
zap.String("protocol_version", serverInfo.ProtocolVersion))
285-
c.upstreamLogger.Debug("[Client→Server] initialize")
286-
if initBytes, err := json.Marshal(initRequest); err == nil {
287-
c.upstreamLogger.Debug(string(initBytes))
288-
}
289-
c.upstreamLogger.Debug("[Server→Client] initialize response")
290-
if respBytes, err := json.Marshal(serverInfo); err == nil {
291-
c.upstreamLogger.Debug(string(respBytes))
327+
328+
// Only log initialization JSON if DEBUG level is enabled
329+
if c.logger.Core().Enabled(zap.DebugLevel) {
330+
c.upstreamLogger.Debug("[Client→Server] initialize")
331+
if initBytes, err := json.Marshal(initRequest); err == nil {
332+
c.upstreamLogger.Debug(string(initBytes))
333+
}
334+
c.upstreamLogger.Debug("[Server→Client] initialize response")
335+
if respBytes, err := json.Marshal(serverInfo); err == nil {
336+
c.upstreamLogger.Debug(string(respBytes))
337+
}
292338
}
293339
}
294340

@@ -437,6 +483,10 @@ func (c *Client) Disconnect() error {
437483

438484
if c.client != nil {
439485
c.logger.Info("Disconnecting from upstream MCP server")
486+
if c.upstreamLogger != nil {
487+
c.upstreamLogger.Info("Disconnecting client")
488+
}
489+
440490
c.client.Close()
441491
c.connected = false
442492
}
@@ -494,23 +544,35 @@ func (c *Client) ListTools(ctx context.Context) ([]*config.ToolMetadata, error)
494544
c.mu.Lock()
495545
c.lastError = err
496546

547+
// Log to both main and server logs for critical errors
548+
c.logger.Error("ListTools failed", zap.Error(err))
549+
if c.upstreamLogger != nil {
550+
c.upstreamLogger.Error("ListTools failed", zap.Error(err))
551+
}
552+
497553
// Check if this is a connection error that indicates the connection is broken
498554
errStr := err.Error()
499555
if strings.Contains(errStr, "broken pipe") ||
500556
strings.Contains(errStr, "connection reset") ||
501557
strings.Contains(errStr, "EOF") ||
502558
strings.Contains(errStr, "connection refused") ||
503559
strings.Contains(errStr, "transport error") {
504-
c.logger.Warn("Connection appears broken, updating state",
505-
zap.String("server", c.config.Name),
506-
zap.Error(err))
560+
561+
// Log pipe errors to both main and server logs
562+
c.logger.Warn("Connection appears broken, updating state", zap.Error(err))
563+
if c.upstreamLogger != nil {
564+
c.upstreamLogger.Warn("Connection broken detected", zap.Error(err))
565+
}
566+
507567
c.connected = false
508568
}
509569
c.mu.Unlock()
510570

511571
return nil, fmt.Errorf("failed to list tools: %w", err)
512572
}
513573

574+
c.logger.Debug("ListTools successful", zap.Int("tools_count", len(toolsResult.Tools)))
575+
514576
// Convert MCP tools to our metadata format
515577
var tools []*config.ToolMetadata
516578
for i := range toolsResult.Tools {
@@ -566,12 +628,16 @@ func (c *Client) CallTool(ctx context.Context, toolName string, args map[string]
566628
zap.String("tool_name", toolName),
567629
zap.Any("args", args))
568630

569-
// Log to upstream logger
631+
// Log detailed transport debug information if DEBUG level is enabled
570632
if c.upstreamLogger != nil {
571633
c.upstreamLogger.Debug("[Client→Server] tools/call",
572634
zap.String("tool", toolName))
573-
if reqBytes, err := json.Marshal(request); err == nil {
574-
c.upstreamLogger.Debug(string(reqBytes))
635+
636+
// Only log full request/response JSON if DEBUG level is enabled
637+
if c.logger.Core().Enabled(zap.DebugLevel) {
638+
if reqBytes, err := json.Marshal(request); err == nil {
639+
c.upstreamLogger.Debug(string(reqBytes))
640+
}
575641
}
576642
}
577643

@@ -580,11 +646,10 @@ func (c *Client) CallTool(ctx context.Context, toolName string, args map[string]
580646
c.mu.Lock()
581647
c.lastError = err
582648

583-
// Log error to upstream logger
649+
// Log to both main and server logs for critical errors
650+
c.logger.Error("CallTool failed", zap.String("tool", toolName), zap.Error(err))
584651
if c.upstreamLogger != nil {
585-
c.upstreamLogger.Error("Tool call failed",
586-
zap.String("tool", toolName),
587-
zap.Error(err))
652+
c.upstreamLogger.Error("Tool call failed", zap.String("tool", toolName), zap.Error(err))
588653
}
589654

590655
// Check if this is a connection error that indicates the connection is broken
@@ -594,26 +659,32 @@ func (c *Client) CallTool(ctx context.Context, toolName string, args map[string]
594659
strings.Contains(errStr, "EOF") ||
595660
strings.Contains(errStr, "connection refused") ||
596661
strings.Contains(errStr, "transport error") {
597-
c.logger.Warn("Connection appears broken during tool call, updating state",
598-
zap.String("server", c.config.Name),
599-
zap.String("tool", toolName),
600-
zap.Error(err))
601-
c.connected = false
602662

663+
// Log pipe errors to both main and server logs
664+
c.logger.Warn("Connection appears broken during tool call, updating state",
665+
zap.String("tool", toolName), zap.Error(err))
603666
if c.upstreamLogger != nil {
604-
c.upstreamLogger.Warn("Connection appears broken during tool call")
667+
c.upstreamLogger.Warn("Connection broken during tool call", zap.Error(err))
605668
}
669+
670+
c.connected = false
606671
}
607672
c.mu.Unlock()
608673

609674
return nil, fmt.Errorf("failed to call tool %s: %w", toolName, err)
610675
}
611676

677+
c.logger.Debug("CallTool successful", zap.String("tool", toolName))
678+
612679
// Log successful response to upstream logger
613680
if c.upstreamLogger != nil {
614681
c.upstreamLogger.Debug("[Server→Client] tools/call response")
615-
if respBytes, err := json.Marshal(result); err == nil {
616-
c.upstreamLogger.Debug(string(respBytes))
682+
683+
// Only log full response JSON if DEBUG level is enabled
684+
if c.logger.Core().Enabled(zap.DebugLevel) {
685+
if respBytes, err := json.Marshal(result); err == nil {
686+
c.upstreamLogger.Debug(string(respBytes))
687+
}
617688
}
618689
}
619690

‎internal/upstream/manager.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -360,8 +360,8 @@ func (m *Manager) GetTotalToolCount() int {
360360
continue
361361
}
362362

363-
// Use a very short timeout to avoid hanging during shutdown
364-
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
363+
// Use a reasonable timeout to allow for Docker container startup and GitHub API calls
364+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
365365
tools, err := client.ListTools(ctx)
366366
cancel()
367367
if err == nil && tools != nil {

0 commit comments

Comments
 (0)