From 251401f9a85a92dfe124dd15e6308916129802ce Mon Sep 17 00:00:00 2001
From: Luiz Kowalski <lhkowalski@gmail.com>
Date: Wed, 31 Jan 2024 16:12:20 -0300
Subject: [PATCH] Jetpack AI: Improve simple site redirect after upgrade
 behavior (#35376)

* Build calypso URL to redirect user back after upgrade

* Changelog

* Only redirect to post edit view if the site is Simple

* Change the default redirect URL for wpcom sites as the calypso home page
---
 ...e-simple-site-redirect-after-upgrade-behavior |  4 ++++
 .../ai-assistant/hooks/use-ai-checkout/index.ts  | 16 +++++++++++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)
 create mode 100644 projects/plugins/jetpack/changelog/update-jetpack-ai-improve-simple-site-redirect-after-upgrade-behavior

diff --git a/projects/plugins/jetpack/changelog/update-jetpack-ai-improve-simple-site-redirect-after-upgrade-behavior b/projects/plugins/jetpack/changelog/update-jetpack-ai-improve-simple-site-redirect-after-upgrade-behavior
new file mode 100644
index 0000000000000..8226d4fff3dc9
--- /dev/null
+++ b/projects/plugins/jetpack/changelog/update-jetpack-ai-improve-simple-site-redirect-after-upgrade-behavior
@@ -0,0 +1,4 @@
+Significance: patch
+Type: other
+
+Jetpack AI: Build a Calypso URL as the redirect destination for Simple sites after upgrade.
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/hooks/use-ai-checkout/index.ts b/projects/plugins/jetpack/extensions/blocks/ai-assistant/hooks/use-ai-checkout/index.ts
index b0c7cb0da932b..598b9aaffc6cc 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/hooks/use-ai-checkout/index.ts
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/hooks/use-ai-checkout/index.ts
@@ -10,6 +10,18 @@ import {
 import useAutosaveAndRedirect from '../../../../shared/use-autosave-and-redirect';
 import useAiFeature from '../use-ai-feature';
 
+const getWPComRedirectToURL = () => {
+	const searchParams = new URLSearchParams( window.location.search );
+	const site = getSiteFragment();
+
+	if ( isSimpleSite() && searchParams.has( 'post' ) ) {
+		// When there is an explicit post, use it as the destination
+		return `https://wordpress.com/post/${ site }/${ searchParams.get( 'post' ) }`;
+	}
+	// When there is no explicit post, or the site is not Simple, use the home page as the destination
+	return `https://wordpress.com/home/${ site }`;
+};
+
 export default function useAICheckout(): {
 	checkoutUrl: string;
 	// eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -18,11 +30,13 @@ export default function useAICheckout(): {
 } {
 	const { nextTier, tierPlansEnabled } = useAiFeature();
 
+	const wpcomRedirectToURL = getWPComRedirectToURL();
+
 	const wpcomCheckoutUrl = tierPlansEnabled
 		? getRedirectUrl( 'jetpack-ai-yearly-tier-upgrade-nudge', {
 				site: getSiteFragment(),
 				path: `jetpack_ai_yearly:-q-${ nextTier?.limit }`,
-				query: `redirect_to=${ encodeURIComponent( window.location.href ) }`,
+				query: `redirect_to=${ encodeURIComponent( wpcomRedirectToURL ) }`,
 		  } )
 		: getRedirectUrl( 'jetpack-ai-monthly-plan-ai-assistant-block-banner', {
 				site: getSiteFragment(),