diff --git a/api/http/handler/websocket/proxy.go b/api/http/handler/websocket/proxy.go
index 9ec317d9b..984240256 100644
--- a/api/http/handler/websocket/proxy.go
+++ b/api/http/handler/websocket/proxy.go
@@ -33,9 +33,12 @@ func (handler *Handler) proxyEdgeAgentWebsocketRequest(w http.ResponseWriter, r
 }
 
 func (handler *Handler) proxyAgentWebsocketRequest(w http.ResponseWriter, r *http.Request, params *webSocketRequestParams) error {
-	// TODO: k8s merge - make sure this is still working with Docker agent
-	//agentURL, err := url.Parse(params.endpoint.URL)
-	agentURL, err := url.Parse(fmt.Sprintf("http://%s", params.endpoint.URL))
+	endpointURL := params.endpoint.URL
+	if params.endpoint.Type == portainer.AgentOnKubernetesEnvironment {
+		endpointURL = fmt.Sprintf("http://%s", params.endpoint.URL)
+	}
+
+	agentURL, err := url.Parse(endpointURL)
 	if err != nil {
 		return err
 	}