From a76191766969c0c4cd98b9fe31005483ad0e422e Mon Sep 17 00:00:00 2001
From: chenhang <averych@126.com>
Date: Tue, 30 Jun 2026 11:14:15 +0800
Subject: [PATCH] fix(chat): return directly to paid image

---
 src/app/chat/chat-screen.tsx                 |  7 -------
 src/app/subscription/subscription-screen.tsx | 22 ++++++++++++++++++--
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/src/app/chat/chat-screen.tsx b/src/app/chat/chat-screen.tsx
index ad89ca87..c68eda23 100644
--- a/src/app/chat/chat-screen.tsx
+++ b/src/app/chat/chat-screen.tsx
@@ -8,7 +8,6 @@ import { useAuthState } from "@/stores/auth/auth-context";
 import { useChatDispatch, useChatState } from "@/stores/chat/chat-context";
 
 import { MobileShell } from "@/app/_components/core";
-import { consumePendingChatImageReturn } from "./chat-image-return-session";
 
 import {
   BrowserHintOverlay,
@@ -89,12 +88,6 @@ export function ChatScreen() {
     authState.loginStatus,
   ]);
 
-  useEffect(() => {
-    const pendingImageReturn = consumePendingChatImageReturn();
-    if (!pendingImageReturn) return;
-    router.replace(pendingImageReturn.returnUrl);
-  }, [router]);
-
   async function handleOpenExternalBrowser(): Promise<void> {
     setShowExternalBrowserDialog(false);
     await openChatInExternalBrowser();
diff --git a/src/app/subscription/subscription-screen.tsx b/src/app/subscription/subscription-screen.tsx
index f5c79d14..d49ed84b 100644
--- a/src/app/subscription/subscription-screen.tsx
+++ b/src/app/subscription/subscription-screen.tsx
@@ -3,6 +3,7 @@ import { useEffect, useMemo, useRef, useState } from "react";
 import { useRouter } from "next/navigation";
 
 import { BackButton } from "@/app/_components";
+import { consumePendingChatImageReturn } from "@/app/chat/chat-image-return-session";
 import { Checkbox, MobileShell } from "@/app/_components/core";
 import { AppConstants } from "@/core/app_constants";
 import { PendingPaymentOrderStorage } from "@/data/storage";
@@ -157,11 +158,28 @@ export function SubscriptionScreen({
     payment.isPollingOrder;
   const canActivate =
     selectedPlan !== null && payment.agreed && !isPaymentBusy;
-  const backHref = returnTo === "chat" ? ROUTES.chat : ROUTES.sidebar;
+
+  const consumeSubscriptionExitUrl = (): string | null => {
+    const pendingImageReturn = consumePendingChatImageReturn();
+    if (pendingImageReturn) return pendingImageReturn.returnUrl;
+    return null;
+  };
+
+  const handleBackClick = () => {
+    router.replace(
+      consumeSubscriptionExitUrl() ??
+        (returnTo === "chat" ? ROUTES.chat : ROUTES.sidebar),
+    );
+  };
 
   const finishPaymentSuccessClose = () => {
     setShowPaymentSuccessDialog(false);
     paymentDispatch({ type: "PaymentReset" });
+    const pendingExitUrl = consumeSubscriptionExitUrl();
+    if (pendingExitUrl) {
+      router.replace(pendingExitUrl);
+      return;
+    }
     if (returnTo === "chat") {
       router.replace(ROUTES.chat);
     }
@@ -219,7 +237,7 @@ export function SubscriptionScreen({
         <header className={styles.header}>
           <BackButton
             className={styles.backSlot}
-            href={backHref}
+            onClick={handleBackClick}
             variant="soft"
           />
         </header>