Merge pull request #346 from rumpl/oauth-only-when-needed

Only start the callback server when it's needed

Author: rumpl

cmd/root/run.go

@@ -23,7 +23,6 @@ import (
 	"github.com/docker/cagent/pkg/chat"
 	"github.com/docker/cagent/pkg/content"
 	"github.com/docker/cagent/pkg/evaluation"
-	"github.com/docker/cagent/pkg/oauth"
 	"github.com/docker/cagent/pkg/remote"
 	"github.com/docker/cagent/pkg/runtime"
 	"github.com/docker/cagent/pkg/session"
@@ -181,31 +180,6 @@ func doRunCommand(ctx context.Context, args []string, exec bool) error {
 			agentFilename = tmpFile.Name()
 		}
 
-		// Set up OAuth redirect URI for CLI/TUI mode
-		if runConfig.RedirectURI == "" {
-			runConfig.RedirectURI = "http://localhost:8083/oauth-callback"
-			slog.Debug("Set default OAuth redirect URI for CLI/TUI mode", "redirectURI", runConfig.RedirectURI)
-
-			// Start OAuth callback server for CLI/TUI mode
-			callbackServer := oauth.NewCallbackServer(8083)
-			err := callbackServer.Start(ctx)
-			if err != nil {
-				slog.Warn("Failed to start OAuth callback server", "error", err)
-			} else {
-				defer func() {
-					shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-					defer cancel()
-					if err := callbackServer.Stop(shutdownCtx); err != nil {
-						slog.Error("Failed to stop OAuth callback server", "error", err)
-					}
-				}()
-				slog.Debug("Started OAuth callback server", "port", 8083)
-
-				// Set up global callback server for OAuth manager
-				oauth.SetGlobalCallbackServer(callbackServer)
-			}
-		}
-
 		agents, err = teamloader.Load(ctx, agentFilename, runConfig)
 		if err != nil {
 			return err

pkg/oauth/interfaces.go

@@ -12,6 +12,9 @@ type Manager interface {
 
 	// SendAuthorizationCode sends the OAuth authorization code after user has completed the OAuth flow
 	SendAuthorizationCode(ctx context.Context, code string) error
+
+	// Cleanup stops any owned resources like callback servers
+	Cleanup(ctx context.Context) error
 }
 
 // ServerInfoProvider interface for toolsets that can provide server information

pkg/oauth/manager.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"sync"
 	"time"
 
 	"github.com/mark3labs/mcp-go/client"
@@ -15,14 +16,48 @@ type manager struct {
 	emitAuthRequired         func(serverURL, serverType, status string)
 	resumeAuthorizeOauthFlow chan bool
 	resumeOauthCodeReceived  chan string
+	callbackServer           *CallbackServer
+	serverMutex              sync.Mutex
+	redirectURI              string
+	port                     int
 }
 
-// NewManager creates a new OAuth manager
-func NewManager(emitAuthRequired func(serverURL, serverType, status string)) Manager {
-	return &manager{
+// NewManager creates a new OAuth manager with optional port configuration
+func NewManager(emitAuthRequired func(serverURL, serverType, status string), opts ...ManagerOption) Manager {
+	m := &manager{
 		emitAuthRequired:         emitAuthRequired,
 		resumeAuthorizeOauthFlow: make(chan bool),
 		resumeOauthCodeReceived:  make(chan string),
+		port:                     8083,
+	}
+
+	// Apply options
+	for _, opt := range opts {
+		opt(m)
+	}
+
+	// Set redirect URI based on port
+	if m.redirectURI == "" {
+		m.redirectURI = fmt.Sprintf("http://localhost:%d/oauth-callback", m.port)
+	}
+
+	return m
+}
+
+// ManagerOption configures the OAuth manager
+type ManagerOption func(*manager)
+
+// WithPort sets the callback server port
+func WithPort(port int) ManagerOption {
+	return func(m *manager) {
+		m.port = port
+	}
+}
+
+// WithRedirectURI sets a custom redirect URI
+func WithRedirectURI(uri string) ManagerOption {
+	return func(m *manager) {
+		m.redirectURI = uri
 	}
 }
 
@@ -136,8 +171,13 @@ func (m *manager) performOAuthAuthorization(ctx context.Context, sessionID strin
 	slog.Debug("Waiting for OAuth authorization code")
 	var code string
 
-	// Check if we have a global callback server running
-	if callbackServer := GetGlobalCallbackServer(); callbackServer != nil {
+	// Ensure callback server is started if needed
+	if err := m.ensureCallbackServer(ctx); err != nil {
+		slog.Warn("Failed to start callback server, falling back to manual input", "error", err)
+	}
+
+	// Check if we have a callback server running (either global or our own)
+	if callbackServer := m.getCallbackServer(); callbackServer != nil {
 		slog.Debug("Using callback server for OAuth authorization")
 		// Wait for callback from the browser
 		callbackCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
@@ -188,3 +228,62 @@ func (m *manager) performOAuthAuthorization(ctx context.Context, sessionID strin
 	slog.Info("OAuth authorization completed successfully")
 	return nil
 }
+
+// ensureCallbackServer starts the callback server if it's not already running
+func (m *manager) ensureCallbackServer(ctx context.Context) error {
+	m.serverMutex.Lock()
+	defer m.serverMutex.Unlock()
+
+	// Check if there's already a global callback server
+	if globalServer := GetGlobalCallbackServer(); globalServer != nil {
+		slog.Debug("Using existing global callback server")
+		return nil
+	}
+
+	// Check if we already have our own server
+	if m.callbackServer != nil {
+		slog.Debug("Callback server already started")
+		return nil
+	}
+
+	// Create and start new callback server
+	slog.Debug("Starting OAuth callback server on demand", "port", m.port)
+	m.callbackServer = NewCallbackServer(m.port)
+	if err := m.callbackServer.Start(ctx); err != nil {
+		return fmt.Errorf("failed to start OAuth callback server: %w", err)
+	}
+
+	// Set as global server so other components can use it
+	SetGlobalCallbackServer(m.callbackServer)
+
+	slog.Debug("OAuth callback server started successfully", "port", m.port)
+	return nil
+}
+
+// getCallbackServer returns the active callback server (either global or local)
+func (m *manager) getCallbackServer() *CallbackServer {
+	// Prefer global server first
+	if globalServer := GetGlobalCallbackServer(); globalServer != nil {
+		return globalServer
+	}
+
+	// Fall back to our local server
+	return m.callbackServer
+}
+
+// Cleanup stops the callback server if we own it
+func (m *manager) Cleanup(ctx context.Context) error {
+	m.serverMutex.Lock()
+	defer m.serverMutex.Unlock()
+
+	if m.callbackServer != nil {
+		slog.Debug("Stopping OAuth callback server")
+		if err := m.callbackServer.Stop(ctx); err != nil {
+			slog.Error("Failed to stop OAuth callback server", "error", err)
+			return err
+		}
+		m.callbackServer = nil
+	}
+
+	return nil
+}

pkg/runtime/runtime.go

@@ -140,6 +140,11 @@ func (r *runtime) handleOAuthAuthorizationFlow(ctx context.Context, sess *sessio
 			events <- AuthorizationRequired(serverURL, serverType, status)
 		}
 		r.oauthManager = oauth.NewManager(emitAuthRequired)
+		defer func() {
+			if cleanupErr := r.oauthManager.Cleanup(ctx); cleanupErr != nil {
+				slog.Error("Failed to cleanup OAuth manager", "error", cleanupErr)
+			}
+		}()
 	}
 
 	return r.oauthManager.HandleAuthorizationFlow(ctx, sess.ID, oauthRequiredErr)