Merge pull request #231 from trungutt/handle-multiple-mcp-oauth

More specificity in toolset start stop calls

Author: dgageot

pkg/agent/agent.go

@@ -5,20 +5,35 @@ import (
 	"fmt"
 	"log/slog"
 	"math/rand"
-	"sync/atomic"
+	"sync"
 
 	"github.com/docker/cagent/pkg/memorymanager"
 	"github.com/docker/cagent/pkg/model/provider"
 	"github.com/docker/cagent/pkg/tools"
 )
 
+// ToolSetStartupError wraps toolset startup failures with context
+type ToolSetStartupError struct {
+	Err   error
+	Index int
+}
+
+func (e *ToolSetStartupError) Error() string {
+	return fmt.Sprintf("failed to start toolset: %v", e.Err)
+}
+
+func (e *ToolSetStartupError) Unwrap() error {
+	return e.Err
+}
+
 // Agent represents an AI agent
 type Agent struct {
 	name               string
 	description        string
 	instruction        string
 	toolsets           []tools.ToolSet
-	toolsetsStarted    atomic.Bool
+	startedToolsets    map[tools.ToolSet]bool
+	toolsetsMutex      sync.RWMutex
 	models             []provider.Provider
 	subAgents          []*Agent
 	parents            []*Agent
@@ -31,8 +46,9 @@ type Agent struct {
 // New creates a new agent
 func New(name, prompt string, opts ...Opt) *Agent {
 	agent := &Agent{
-		name:        name,
-		instruction: prompt,
+		name:            name,
+		instruction:     prompt,
+		startedToolsets: make(map[tools.ToolSet]bool),
 	}
 
 	for _, opt := range opts {
@@ -129,31 +145,46 @@ func (a *Agent) ToolSets() []tools.ToolSet {
 }
 
 func (a *Agent) ensureToolSetsAreStarted(ctx context.Context) error {
-	if a.toolsetsStarted.Load() {
-		return nil
-	}
+	a.toolsetsMutex.Lock()
+	defer a.toolsetsMutex.Unlock()
+
+	for i, toolSet := range a.toolsets {
+		// Skip if toolset is already started
+		if a.startedToolsets[toolSet] {
+			continue
+		}
 
-	for _, toolSet := range a.toolsets {
 		if err := toolSet.Start(ctx); err != nil {
-			return fmt.Errorf("failed to start toolset: %w", err)
+			return &ToolSetStartupError{
+				Err:   err,
+				Index: i,
+			}
 		}
+
+		// Mark toolset as started
+		a.startedToolsets[toolSet] = true
 	}
 
-	a.toolsetsStarted.Store(true)
 	return nil
 }
 
 func (a *Agent) StopToolSets() error {
-	if !a.toolsetsStarted.Load() {
-		return nil
-	}
+	a.toolsetsMutex.Lock()
+	defer a.toolsetsMutex.Unlock()
 
 	for _, toolSet := range a.toolsets {
+		// Only stop toolsets that are marked as started
+		if !a.startedToolsets[toolSet] {
+			continue
+		}
+
 		if err := toolSet.Stop(); err != nil {
 			return fmt.Errorf("failed to stop toolset: %w", err)
 		}
+
+		// Mark toolset as stopped
+		a.startedToolsets[toolSet] = false
 	}
 
-	a.toolsetsStarted.Store(false)
 	return nil
 }

pkg/runtime/runtime.go

@@ -126,12 +126,18 @@ func (r *runtime) registerDefaultTools() {
 func (r *runtime) getAgentToolsWithOAuthHandling(ctx context.Context, a *agent.Agent) ([]tools.Tool, error) {
 	agentTools, err := a.Tools(ctx)
 	if err != nil {
-		// If this is an OAuth authorization error, wrap it with server info
-		if client.IsOAuthAuthorizationRequiredError(err) {
-			// Try to find which toolset caused the OAuth error by checking each one
-			for _, toolSet := range a.ToolSets() {
-				if serverInfoProvider, ok := toolSet.(oauth.ServerInfoProvider); ok {
-					return nil, oauth.WrapOAuthError(err, serverInfoProvider)
+		// Check if this is a ToolSetStartupError
+		var toolSetStartupErr *agent.ToolSetStartupError
+		if errors.As(err, &toolSetStartupErr) {
+			// Check if the inner error is an OAuth authorization error
+			if client.IsOAuthAuthorizationRequiredError(toolSetStartupErr.Err) {
+				// Use the index from ToolSetStartupError to get the specific toolset
+				toolsets := a.ToolSets()
+				if toolSetStartupErr.Index >= 0 && toolSetStartupErr.Index < len(toolsets) {
+					toolSet := toolsets[toolSetStartupErr.Index]
+					if serverInfoProvider, ok := toolSet.(oauth.ServerInfoProvider); ok {
+						return nil, oauth.WrapOAuthError(toolSetStartupErr.Err, serverInfoProvider)
+					}
 				}
 			}
 		}

pkg/tools/mcp/toolset.go

@@ -74,24 +74,26 @@ func (t *Toolset) Tools(ctx context.Context) ([]tools.Tool, error) {
 
 // Start starts the toolset
 func (t *Toolset) Start(ctx context.Context) error {
-	slog.Debug("Starting MCP toolset")
+	serverURL, _ := t.c.GetServerInfo()
+	slog.Debug("Starting MCP toolset", "server", serverURL)
 	err := t.c.Start(ctx)
 	if err != nil {
-		slog.Error("Failed to start MCP toolset", "error", err)
+		slog.Error("Failed to start MCP toolset", "server", serverURL, "error", err)
 		return err
 	}
-	slog.Debug("Started MCP toolset successfully")
+	slog.Debug("Started MCP toolset successfully", "server", serverURL)
 	return nil
 }
 
 // Stop stops the toolset
 func (t *Toolset) Stop() error {
-	slog.Debug("Stopping MCP toolset")
+	serverURL, _ := t.c.GetServerInfo()
+	slog.Debug("Stopping MCP toolset", "server", serverURL)
 	err := t.c.Stop()
 	if err != nil {
-		slog.Error("Failed to stop MCP toolset", "error", err)
+		slog.Error("Failed to stop MCP toolset", "server", serverURL, "error", err)
 		return err
 	}
-	slog.Debug("Stopped MCP toolset successfully")
+	slog.Debug("Stopped MCP toolset successfully", "server", serverURL)
 	return nil
 }