diff --git a/.psalm/stubs.php b/.psalm/stubs.php
index 27484ec28..0ac41fcd3 100644
--- a/.psalm/stubs.php
+++ b/.psalm/stubs.php
@@ -153,3 +153,26 @@ class WP_HTML_Tag_Processor {
 		return '';
 	}
 }
+
+/**
+ * WooCommerce Session stubs for Psalm
+ */
+class WC_Session {
+	/**
+	 * Get session cookie.
+	 *
+	 * @return string|false
+	 */
+	public function get_session_cookie() {
+		return '';
+	}
+
+	/**
+	 * Set customer session cookie.
+	 *
+	 * @param bool $set Whether to set the cookie.
+	 * @return void
+	 */
+	public function set_customer_session_cookie( $set ) {
+	}
+}
diff --git a/modules/ppcp-api-client/src/Entity/Order.php b/modules/ppcp-api-client/src/Entity/Order.php
index 7a8a596e0..636323fab 100644
--- a/modules/ppcp-api-client/src/Entity/Order.php
+++ b/modules/ppcp-api-client/src/Entity/Order.php
@@ -185,6 +185,11 @@ class Order {
 		return $this->payment_source;
 	}
 
+	/**
+	 * Returns the links.
+	 *
+	 * @return mixed|null
+	 */
 	public function links() {
 		return $this->links;
 	}
diff --git a/modules/ppcp-wc-gateway/src/Endpoint/ReturnUrlEndpoint.php b/modules/ppcp-wc-gateway/src/Endpoint/ReturnUrlEndpoint.php
index 91900daf9..96d389613 100644
--- a/modules/ppcp-wc-gateway/src/Endpoint/ReturnUrlEndpoint.php
+++ b/modules/ppcp-wc-gateway/src/Endpoint/ReturnUrlEndpoint.php
@@ -83,7 +83,8 @@ class ReturnUrlEndpoint {
 			WC()->initialize_session();
 		}
 
-		if ( ! WC()->session->get_session_cookie() ) {
+		// Set customer session cookie if not already set.
+		if ( WC()->session && ! WC()->session->get_session_cookie() ) {
 			WC()->session->set_customer_session_cookie( true );
 		}
 
@@ -98,7 +99,7 @@ class ReturnUrlEndpoint {
 				$wc_order_id = (int) $order->purchase_units()[0]->custom_id();
 				$wc_order    = wc_get_order( $wc_order_id );
 
-				if ( ! $wc_order ) {
+				if ( ! $wc_order || ! is_a( $wc_order, \WC_Order::class ) ) {
 					wc_add_notice( __( 'Order information is missing. Please try placing your order again.', 'woocommerce-paypal-payments' ), 'error' );
 					wp_safe_redirect( wc_get_checkout_url() );
 					exit();