refactor(oauth): scope callback serving to authorization

cmd/root/debug_oauth_test.go

@@ -37,7 +37,7 @@ func TestFindMCPRemote_MatchesTopLevelMCPByName(t *testing.T) {
 // TestFindMCPRemote_ForwardsCallbackPortAndRedirectURL proves the full
 // RemoteOAuthConfig — including CallbackPort and CallbackRedirectURL —
 // survives findMCPRemote's lookup verbatim, so PerformOAuthLogin's
-// runtime-aligned callback wiring (NewCallbackServerOnPort/
+// runtime-aligned callback wiring (RequestAuthorizationCode/
 // ResolveRedirectURI) sees the exact same per-toolset config the runtime
 // would use for the same remote MCP server.
 func TestFindMCPRemote_ForwardsCallbackPortAndRedirectURL(t *testing.T) {

pkg/runtime/remote_oauth_callback_test.go

@@ -0,0 +1,100 @@
+package runtime
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/tools"
+	"github.com/docker/docker-agent/pkg/tools/mcp/oauthflow"
+)
+
+func TestRemoteOAuthCallbackFlow(t *testing.T) {
+	for _, failBrowser := range []bool{false, true} {
+		t.Run(map[bool]string{false: "success", true: "browser failure"}[failBrowser], func(t *testing.T) {
+			// OAuth dependency overrides are process-wide; do not parallelize.
+			var registeredRedirect, exchangedRedirect string
+			var exchanged bool
+			mux := http.NewServeMux()
+			server := httptest.NewServer(mux)
+			defer server.Close()
+			mux.HandleFunc("/register", func(w http.ResponseWriter, _ *http.Request) {
+				w.Header().Set("Content-Type", "application/json")
+				_, _ = io.WriteString(w, `{"client_id":"registered"}`)
+			})
+			mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) {
+				_ = r.ParseForm()
+				exchangedRedirect = r.FormValue("redirect_uri")
+				exchanged = true
+				w.Header().Set("Content-Type", "application/json")
+				_, _ = io.WriteString(w, `{"access_token":"token","token_type":"Bearer"}`)
+			})
+			restoreClient := oauthflow.SetHTTPClientForTesting(server.Client())
+			defer restoreClient()
+			failure := errors.New("cannot open browser")
+			restoreBrowser := oauthflow.SetBrowserOpenerForTesting(func(ctx context.Context, authURL string) error {
+				if failBrowser {
+					return failure
+				}
+				parsed, err := url.Parse(authURL)
+				if err != nil {
+					return err
+				}
+				registeredRedirect = parsed.Query().Get("redirect_uri")
+				callback, err := url.Parse(registeredRedirect)
+				if err != nil {
+					return err
+				}
+				query := callback.Query()
+				query.Set("code", "code")
+				query.Set("state", parsed.Query().Get("state"))
+				callback.RawQuery = query.Encode()
+				request, err := http.NewRequestWithContext(ctx, http.MethodGet, callback.String(), http.NoBody)
+				if err != nil {
+					return err
+				}
+				response, err := server.Client().Do(request)
+				if err != nil {
+					return err
+				}
+				defer response.Body.Close()
+				_, err = io.Copy(io.Discard, response.Body)
+				return err
+			})
+			defer restoreBrowser()
+			client := &elicitationRecordingClient{}
+			rt, err := NewRemoteRuntime(client)
+			require.NoError(t, err)
+			err = rt.handleOAuthElicitation(t.Context(), "session", &ElicitationRequestEvent{
+				ElicitationID: "oauth-request",
+				Meta: map[string]any{
+					"docker-agent/server_url": server.URL,
+					"auth_server_metadata": map[string]any{
+						"issuer":                 server.URL,
+						"authorization_endpoint": server.URL + "/authorize",
+						"registration_endpoint":  server.URL + "/register",
+						"token_endpoint":         server.URL + "/token",
+					},
+				},
+			})
+			if failBrowser {
+				require.ErrorIs(t, err, failure)
+				assert.False(t, exchanged)
+				assert.Equal(t, []tools.ElicitationAction{tools.ElicitationActionDecline}, client.actions)
+			} else {
+				require.NoError(t, err)
+				assert.True(t, exchanged, "exchange must use the still-live parent context")
+				assert.Equal(t, registeredRedirect, exchangedRedirect)
+				assert.Equal(t, []tools.ElicitationAction{tools.ElicitationActionAccept}, client.actions)
+			}
+			assert.Equal(t, []string{"oauth-request"}, client.ids)
+		})
+	}
+}

pkg/runtime/remote_runtime.go

@@ -747,79 +747,58 @@ func (r *RemoteRuntime) handleOAuthElicitation(ctx context.Context, sessionID st
 	oauthCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
 	defer cancel()
 
-	slog.DebugContext(ctx, "Creating OAuth callback server")
-	callbackServer := oauthflow.NewCallbackServer()
-	defer func() {
-		// Detach from ctx's cancellation (the request may be done) but
-		// keep its trace context for the shutdown.
-		shutdownCtx, shutdownCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
-		defer shutdownCancel()
-		if err := callbackServer.Shutdown(shutdownCtx); err != nil {
-			slog.ErrorContext(ctx, "Failed to shutdown callback server", "error", err)
+	var redirectURI, clientID, clientSecret, authMethod, verifier, state string
+	prepared := false
+	code, receivedState, err := oauthflow.RequestAuthorizationCode(oauthCtx, 0, func(ctx context.Context, address oauthflow.CallbackAddress) (oauthflow.AuthorizationRequest, error) {
+		redirectURI = address.RedirectURI()
+		if authMetadata.RegistrationEndpoint != "" {
+			slog.DebugContext(ctx, "Attempting dynamic client registration")
+			clientID, clientSecret, authMethod, err = oauthflow.RegisterClient(ctx, &authMetadata, redirectURI, nil)
+			if err != nil {
+				slog.ErrorContext(ctx, "Dynamic client registration failed", "error", err)
+				return oauthflow.AuthorizationRequest{}, fmt.Errorf("failed to register client: %w", err)
+			}
+			slog.DebugContext(ctx, "Client registered successfully", "client_id", clientID)
+		} else {
+			err := errors.New("authorization server does not support dynamic client registration")
+			slog.ErrorContext(ctx, "Client registration not supported", "error", err)
+			return oauthflow.AuthorizationRequest{}, err
 		}
-	}()
-
-	if err := callbackServer.Start(ctx); err != nil {
-		slog.ErrorContext(ctx, "Failed to start callback server", "error", err)
-		_ = r.client.ResumeElicitation(ctx, sessionID, "decline", nil, req.ElicitationID)
-		return fmt.Errorf("failed to start callback server: %w", err)
-	}
-
-	redirectURI := callbackServer.GetRedirectURI()
-	slog.DebugContext(ctx, "Callback server started", "redirect_uri", redirectURI)
 
-	var clientID, clientSecret, authMethod string
-	if authMetadata.RegistrationEndpoint != "" {
-		slog.DebugContext(ctx, "Attempting dynamic client registration")
-		clientID, clientSecret, authMethod, err = oauthflow.RegisterClient(oauthCtx, &authMetadata, redirectURI, nil)
+		state, err = oauthflow.GenerateState()
 		if err != nil {
-			slog.ErrorContext(ctx, "Dynamic client registration failed", "error", err)
-			_ = r.client.ResumeElicitation(ctx, sessionID, "decline", nil, req.ElicitationID)
-			return fmt.Errorf("failed to register client: %w", err)
+			slog.ErrorContext(ctx, "Failed to generate state", "error", err)
+			return oauthflow.AuthorizationRequest{}, fmt.Errorf("failed to generate state: %w", err)
 		}
-		slog.DebugContext(ctx, "Client registered successfully", "client_id", clientID)
-	} else {
-		err := errors.New("authorization server does not support dynamic client registration")
-		slog.ErrorContext(ctx, "Client registration not supported", "error", err)
-		_ = r.client.ResumeElicitation(ctx, sessionID, "decline", nil, req.ElicitationID)
-		return err
-	}
 
-	state, err := oauthflow.GenerateState()
-	if err != nil {
-		slog.ErrorContext(ctx, "Failed to generate state", "error", err)
-		_ = r.client.ResumeElicitation(ctx, sessionID, "decline", nil, req.ElicitationID)
-		return fmt.Errorf("failed to generate state: %w", err)
-	}
+		verifier = oauthflow.GeneratePKCEVerifier()
 
-	callbackServer.SetExpectedState(state)
-	verifier := oauthflow.GeneratePKCEVerifier()
-
-	authURL := oauthflow.BuildAuthorizationURL(
-		authMetadata.AuthorizationEndpoint,
-		clientID,
-		redirectURI,
-		state,
-		oauth2.S256ChallengeFromVerifier(verifier),
-		resourceIndicator,
-		nil,
-	)
+		authURL := oauthflow.BuildAuthorizationURL(
+			authMetadata.AuthorizationEndpoint,
+			clientID,
+			redirectURI,
+			state,
+			oauth2.S256ChallengeFromVerifier(verifier),
+			resourceIndicator,
+			nil,
+		)
 
-	slog.DebugContext(ctx, "Authorization URL built", "url", authURL)
+		slog.DebugContext(ctx, "Authorization URL built", "url", authURL)
 
-	slog.DebugContext(ctx, "Requesting authorization code")
-	code, receivedState, err := oauthflow.RequestAuthorizationCode(oauthCtx, authURL, callbackServer, state)
+		prepared = true
+		return oauthflow.AuthorizationRequest{URL: authURL, State: state}, nil
+	})
 	if err != nil {
 		slog.ErrorContext(ctx, "Failed to get authorization code", "error", err)
 		_ = r.client.ResumeElicitation(ctx, sessionID, "decline", nil, req.ElicitationID)
+		if !prepared {
+			return err
+		}
 		return fmt.Errorf("failed to get authorization code: %w", err)
 	}
-
 	if receivedState != state {
-		err := fmt.Errorf("state mismatch: expected %s, got %s", state, receivedState)
-		slog.ErrorContext(ctx, "State mismatch in authorization response", "error", err)
 		_ = r.client.ResumeElicitation(ctx, sessionID, "decline", nil, req.ElicitationID)
-		return err
+		return errors.New("state mismatch in authorization response")
 	}
 
 	slog.DebugContext(ctx, "Authorization code received, exchanging for token")

pkg/tools/mcp/oauth.go

@@ -457,7 +457,7 @@ func callbackRedirectURLFrom(c *latest.RemoteOAuthConfig) string {
 
 // callbackPortFrom is a nil-safe accessor for the optional CallbackPort
 // field on a RemoteOAuthConfig; zero means "let the OS pick a free port",
-// both here and as NewCallbackServerOnPort's port argument.
+// both here and for the callback listener.
 func callbackPortFrom(c *latest.RemoteOAuthConfig) int {
 	if c == nil {
 		return 0
@@ -1270,74 +1270,66 @@ func (t *oauthTransport) handleManagedOAuthFlow(ctx context.Context, authServer,
 		return fmt.Errorf("failed to fetch authorization server metadata: %w", err)
 	}
 
-	slog.DebugContext(ctx, "Starting OAuth callback server")
-	callbackServer, err := startCallbackServer(ctx, callbackPortFrom(t.oauthConfig))
-	if err != nil {
-		return err
-	}
-	defer func() {
-		// Detach from ctx's cancellation (the request may be done) but
-		// keep its trace context for the shutdown.
-		shutdownCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
-		defer cancel()
-		if err := callbackServer.Shutdown(shutdownCtx); err != nil {
-			slog.ErrorContext(ctx, "Failed to shutdown callback server", "error", err)
-		}
-	}()
-
-	redirectURI := callbackServer.ResolveRedirectURI(callbackRedirectURLFrom(t.oauthConfig))
-	slog.DebugContext(ctx, "Using redirect URI", "uri", redirectURI)
+	var redirectURI, clientID, clientSecret, authMethod, verifier, state string
+	var scopes []string
+	resourceIndicator := cmp.Or(resourceMetadata.Resource, t.baseURL)
+	prepared := false
+	code, receivedState, err := RequestAuthorizationCode(ctx, callbackPortFrom(t.oauthConfig), func(ctx context.Context, address CallbackAddress) (AuthorizationRequest, error) {
+		redirectURI = address.ResolveRedirectURI(callbackRedirectURLFrom(t.oauthConfig))
+		slog.DebugContext(ctx, "Using redirect URI", "uri", redirectURI)
 
-	clientID, clientSecret, scopes, authMethod, err := t.resolveClientCredentials(ctx, authServerMetadata, redirectURI, challengeScopesFromWWWAuth(wwwAuth), resourceMetadata.ScopesSupported)
-	if err != nil {
-		return err
-	}
+		clientID, clientSecret, scopes, authMethod, err = t.resolveClientCredentials(ctx, authServerMetadata, redirectURI, challengeScopesFromWWWAuth(wwwAuth), resourceMetadata.ScopesSupported)
+		if err != nil {
+			return AuthorizationRequest{}, err
+		}
 
-	state, err := GenerateState()
-	if err != nil {
-		return fmt.Errorf("failed to generate state: %w", err)
-	}
+		state, err = GenerateState()
+		if err != nil {
+			return AuthorizationRequest{}, fmt.Errorf("failed to generate state: %w", err)
+		}
 
-	callbackServer.SetExpectedState(state)
-	verifier := GeneratePKCEVerifier()
-	resourceIndicator := cmp.Or(resourceMetadata.Resource, t.baseURL)
+		verifier = GeneratePKCEVerifier()
 
-	authURL := BuildAuthorizationURL(
-		authServerMetadata.AuthorizationEndpoint,
-		clientID,
-		redirectURI,
-		state,
-		oauth2.S256ChallengeFromVerifier(verifier),
-		resourceIndicator,
-		scopes,
-	)
+		authURL := BuildAuthorizationURL(
+			authServerMetadata.AuthorizationEndpoint,
+			clientID,
+			redirectURI,
+			state,
+			oauth2.S256ChallengeFromVerifier(verifier),
+			resourceIndicator,
+			scopes,
+		)
 
-	result, err := t.elicit(ctx, &mcpsdk.ElicitParams{
-		Message:         fmt.Sprintf("The MCP server at %s requires OAuth authorization. Do you want to proceed?", t.baseURL),
-		RequestedSchema: nil,
-		Meta: map[string]any{
-			"docker-agent/type":       "oauth_flow",
-			"docker-agent/server_url": t.baseURL,
-		},
-	})
-	if err != nil {
-		return fmt.Errorf("failed to send elicitation request: %w", err)
-	}
+		result, err := t.elicit(ctx, &mcpsdk.ElicitParams{
+			Message:         fmt.Sprintf("The MCP server at %s requires OAuth authorization. Do you want to proceed?", t.baseURL),
+			RequestedSchema: nil,
+			Meta: map[string]any{
+				"docker-agent/type":       "oauth_flow",
+				"docker-agent/server_url": t.baseURL,
+			},
+		})
+		if err != nil {
+			return AuthorizationRequest{}, fmt.Errorf("failed to send elicitation request: %w", err)
+		}
 
-	slog.DebugContext(ctx, "Elicitation response received", "result", result)
+		slog.DebugContext(ctx, "Elicitation response received", "result", result)
 
-	if result.Action != tools.ElicitationActionAccept {
-		return errors.New("user declined OAuth authorization")
-	}
+		if result.Action != tools.ElicitationActionAccept {
+			return AuthorizationRequest{}, errors.New("user declined OAuth authorization")
+		}
 
-	slog.DebugContext(ctx, "Requesting authorization code", "url", sanitizeURLForLog(authURL))
-	span.AddEvent("oauth.step", trace.WithAttributes(attribute.String("cagent.oauth.step", "request_authorization_code")))
+		slog.DebugContext(ctx, "Requesting authorization code", "url", sanitizeURLForLog(authURL))
+		span.AddEvent("oauth.step", trace.WithAttributes(attribute.String("cagent.oauth.step", "request_authorization_code")))
 
-	code, receivedState, err := RequestAuthorizationCode(ctx, authURL, callbackServer, state)
+		prepared = true
+		return AuthorizationRequest{URL: authURL, State: state}, nil
+	})
 	if err != nil {
+		if !prepared {
+			return err
+		}
 		return fmt.Errorf("failed to get authorization code: %w", err)
 	}
-
 	if receivedState != state {
 		return errors.New("state mismatch in authorization response")
 	}