fix(private-room): preserve payment return route

src/app/private-room/__tests__/use-private-room-flow-navigation.test.tsx

@@ -0,0 +1,112 @@
+import { act, type Dispatch } from "react";
+import { createRoot, type Root } from "react-dom/client";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+import type { LoginStatus } from "@/data/dto/auth";
+import { ROUTES } from "@/router/routes";
+import type { PrivateRoomEvent } from "@/stores/private-room";
+import type { PrivateRoomUnlockPaywallRequest } from "@/stores/private-room/private-room-state";
+
+import { usePrivateRoomUnlockPaywallNavigation } from "../use-private-room-flow";
+
+const mocks = vi.hoisted(() => ({
+  openAuth: vi.fn(),
+  openSubscription: vi.fn(),
+  paywallShown: vi.fn(),
+}));
+
+vi.mock("@/router/use-app-navigator", () => ({
+  useAppNavigator: () => ({
+    openAuth: mocks.openAuth,
+    openSubscription: mocks.openSubscription,
+  }),
+}));
+
+vi.mock("@/lib/analytics", () => ({
+  behaviorAnalytics: {
+    paywallShown: mocks.paywallShown,
+  },
+}));
+
+const unlockPaywallRequest: PrivateRoomUnlockPaywallRequest = {
+  albumId: "album-1",
+  reason: "insufficient_credits",
+  requiredCredits: 10,
+  currentCredits: 3,
+  shortfallCredits: 7,
+};
+
+describe("Private Room paywall navigation", () => {
+  let container: HTMLDivElement;
+  let root: Root;
+
+  beforeEach(() => {
+    (globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean })
+      .IS_REACT_ACT_ENVIRONMENT = true;
+    mocks.openAuth.mockReset();
+    mocks.openSubscription.mockReset();
+    mocks.paywallShown.mockReset();
+    container = document.createElement("div");
+    document.body.appendChild(container);
+    root = createRoot(container);
+  });
+
+  afterEach(() => {
+    act(() => root.unmount());
+    container.remove();
+  });
+
+  it("opens top up with a private room return target", () => {
+    const roomDispatch = vi.fn<Dispatch<PrivateRoomEvent>>();
+
+    renderHarness(root, "email", roomDispatch);
+
+    expect(mocks.openSubscription).toHaveBeenCalledWith({
+      type: "topup",
+      returnTo: "private-room",
+      analytics: {
+        entryPoint: "private_album_unlock",
+        triggerReason: "insufficient_credits",
+      },
+    });
+    expect(roomDispatch).toHaveBeenCalledWith({
+      type: "PrivateRoomUnlockPaywallConsumed",
+    });
+  });
+
+  it("keeps guest users on the private room auth return path", () => {
+    const roomDispatch = vi.fn<Dispatch<PrivateRoomEvent>>();
+
+    renderHarness(root, "guest", roomDispatch);
+
+    expect(mocks.openAuth).toHaveBeenCalledWith(ROUTES.privateRoom);
+    expect(mocks.openSubscription).not.toHaveBeenCalled();
+  });
+});
+
+function renderHarness(
+  root: Root,
+  loginStatus: LoginStatus,
+  roomDispatch: Dispatch<PrivateRoomEvent>,
+): void {
+  act(() => {
+    root.render(
+      <Harness loginStatus={loginStatus} roomDispatch={roomDispatch} />,
+    );
+  });
+}
+
+function Harness({
+  loginStatus,
+  roomDispatch,
+}: {
+  loginStatus: LoginStatus;
+  roomDispatch: Dispatch<PrivateRoomEvent>;
+}) {
+  usePrivateRoomUnlockPaywallNavigation({
+    loginStatus,
+    roomDispatch,
+    unlockPaywallRequest,
+  });
+  return null;
+}

src/app/private-room/private-room-screen.tsx

@@ -111,6 +111,7 @@ export function PrivateRoomScreen() {
     }
     navigator.openSubscription({
       type: "topup",
+      returnTo: "private-room",
       analytics: {
         entryPoint: "private_room",
         triggerReason: "vip_cta",

src/app/private-room/use-private-room-flow.ts

@@ -91,6 +91,7 @@ export function usePrivateRoomUnlockPaywallNavigation({
       });
       navigator.openSubscription({
         type: "topup",
+        returnTo: "private-room",
         analytics: {
           entryPoint: "private_album_unlock",
           triggerReason: "insufficient_credits",

src/app/subscription/components/subscription-checkout-button.tsx

@@ -6,6 +6,7 @@
  */
 import { usePaymentLaunchFlow } from "@/app/_hooks/use-payment-launch-flow";
 import { PaymentLaunchDialogs } from "@/app/_components/payment/payment-launch-dialogs";
+import type { SubscriptionReturnTo } from "@/lib/navigation/subscription_exit";
 import {
   usePaymentDispatch,
   usePaymentState,
@@ -20,7 +21,7 @@ export interface SubscriptionCheckoutButtonProps {
   /** 是否可用（未选套餐 / 未勾选协议 → false） */
   disabled?: boolean;
   subscriptionType: "vip" | "topup";
-  returnTo?: "chat" | null;
+  returnTo?: SubscriptionReturnTo;
 }
 
 export function SubscriptionCheckoutButton({

src/app/subscription/page.tsx

@@ -6,6 +6,7 @@ import {
 import {
   getFirstPaymentSearchParam,
   parsePaymentReturnSearchParams,
+  parseSubscriptionReturnTo,
   type PaymentSearchParams,
 } from "@/lib/payment/payment_search_params";
 
@@ -38,7 +39,7 @@ export default async function SubscriptionPage({
     <SubscriptionScreen
       subscriptionType={subscriptionType}
       shouldResumePendingOrder={paymentReturn.shouldResumePendingOrder}
-      returnTo={toReturnTo(getFirstPaymentSearchParam(query.returnTo))}
+      returnTo={parseSubscriptionReturnTo(query.returnTo)}
       initialPayChannel={paymentReturn.initialPayChannel}
       analyticsContext={analyticsContext}
     />
@@ -48,7 +49,3 @@ export default async function SubscriptionPage({
 function toSubscriptionType(value: string | null): SubscriptionType {
   return value === "topup" ? "topup" : "vip";
 }
-
-function toReturnTo(value: string | null): "chat" | null {
-  return value === "chat" ? "chat" : null;
-}

src/app/subscription/subscription-screen.tsx

@@ -6,6 +6,7 @@ import { Checkbox, MobileShell } from "@/app/_components/core";
 import { usePaymentPlanAnalytics } from "@/app/_hooks/use-payment-plan-analytics";
 import { AppConstants } from "@/core/app_constants";
 import type { PayChannel } from "@/data/dto/payment";
+import type { SubscriptionReturnTo } from "@/lib/navigation/subscription_exit";
 import {
   behaviorAnalytics,
   getDefaultPaymentAnalyticsContext,
@@ -38,7 +39,7 @@ export type { SubscriptionType } from "./subscription-screen.helpers";
 export interface SubscriptionScreenProps {
   subscriptionType?: SubscriptionType;
   shouldResumePendingOrder?: boolean;
-  returnTo?: "chat" | null;
+  returnTo?: SubscriptionReturnTo;
   initialPayChannel?: PayChannel | null;
   analyticsContext?: PaymentAnalyticsContext;
 }

src/app/subscription/use-subscription-payment-flow.ts

@@ -5,13 +5,14 @@ import { useEffect, useRef, useState } from "react";
 import { usePaymentRouteFlow } from "@/app/_hooks/use-payment-route-flow";
 import { useAppNavigator } from "@/router/use-app-navigator";
 import type { PayChannel } from "@/data/dto/payment";
+import type { SubscriptionReturnTo } from "@/lib/navigation/subscription_exit";
 
 import type { SubscriptionType } from "./subscription-screen.helpers";
 
 export interface UseSubscriptionPaymentFlowInput {
   subscriptionType: SubscriptionType;
   shouldResumePendingOrder: boolean;
-  returnTo: "chat" | null;
+  returnTo: SubscriptionReturnTo;
   initialPayChannel: PayChannel;
 }
 

src/data/storage/payment/pending_payment_order_storage.ts

@@ -12,7 +12,7 @@ const PendingPaymentOrderSchema = z.object({
   payChannel: z.literal("ezpay"),
   subscriptionType: z.enum(["vip", "topup", "tip"]),
   tipCoffeeType: z.enum(["small", "medium", "large"]).optional(),
-  returnTo: z.enum(["chat"]).optional(),
+  returnTo: z.enum(["chat", "private-room"]).optional(),
   createdAt: z.number(),
 });
 

src/lib/navigation/__tests__/subscription_exit.test.ts

@@ -12,6 +12,7 @@ import {
   consumeSubscriptionExplicitExitUrl,
   getSubscriptionFallbackExitUrl,
   peekSubscriptionExplicitExitUrl,
+  resolveSubscriptionSuccessExitUrl,
 } from "../subscription_exit";
 
 vi.mock("../chat_image_return_session", () => ({
@@ -65,6 +66,69 @@ describe("subscription exit helpers", () => {
     );
   });
 
+  it("returns directly to private room without consuming chat state", async () => {
+    consumePendingChatImageReturnMock.mockResolvedValue({
+      reason: "image_paywall",
+      messageId: "msg_1",
+      returnUrl: "/chat?image=msg_1",
+      createdAt: 1,
+    });
+    peekPendingChatUnlockMock.mockResolvedValue({
+      reason: "single_message_unlock",
+      displayMessageId: "msg_1",
+      messageId: "msg_1",
+      kind: "image",
+      returnUrl: "/chat?image=msg_1",
+      stage: "payment",
+      createdAt: 1,
+    });
+
+    expect(getSubscriptionFallbackExitUrl("private-room")).toBe(
+      ROUTES.privateRoom,
+    );
+    await expect(consumeSubscriptionExitUrl("private-room")).resolves.toBe(
+      ROUTES.privateRoom,
+    );
+    expect(consumePendingChatImageReturnMock).not.toHaveBeenCalled();
+    expect(peekPendingChatUnlockMock).not.toHaveBeenCalled();
+    expect(clearPendingChatUnlockMock).not.toHaveBeenCalled();
+  });
+
+  it("keeps private room ahead of stale chat state after payment", async () => {
+    peekPendingChatUnlockMock.mockResolvedValue({
+      reason: "single_message_unlock",
+      displayMessageId: "msg_1",
+      messageId: "msg_1",
+      kind: "image",
+      returnUrl: "/chat?image=msg_1",
+      stage: "payment",
+      createdAt: 1,
+    });
+
+    await expect(
+      resolveSubscriptionSuccessExitUrl("private-room"),
+    ).resolves.toBe(ROUTES.privateRoom);
+    expect(peekPendingChatUnlockMock).not.toHaveBeenCalled();
+    expect(clearPendingChatUnlockMock).not.toHaveBeenCalled();
+  });
+
+  it("keeps pending chat unlocks ahead of chat fallback after payment", async () => {
+    peekPendingChatUnlockMock.mockResolvedValue({
+      reason: "single_message_unlock",
+      displayMessageId: "msg_1",
+      messageId: "msg_1",
+      kind: "image",
+      returnUrl: "/chat?image=msg_1",
+      stage: "payment",
+      createdAt: 1,
+    });
+
+    await expect(resolveSubscriptionSuccessExitUrl("chat")).resolves.toBe(
+      "/chat?image=msg_1",
+    );
+    expect(clearPendingChatUnlockMock).not.toHaveBeenCalled();
+  });
+
   it("peeks chat unlock return urls without clearing them", async () => {
     peekPendingChatUnlockMock.mockResolvedValue({
       reason: "single_message_unlock",

src/lib/navigation/subscription_exit.ts

@@ -1,6 +1,7 @@
 "use client";
 
 import { ROUTES } from "@/router/routes";
+import type { AppSubscriptionReturnTo } from "@/router/navigation-types";
 
 import { consumePendingChatImageReturn } from "./chat_image_return_session";
 import {
@@ -8,7 +9,7 @@ import {
   peekPendingChatUnlock,
 } from "./chat_unlock_session";
 
-export type SubscriptionReturnTo = "chat" | null;
+export type SubscriptionReturnTo = AppSubscriptionReturnTo;
 
 export async function consumeSubscriptionExplicitExitUrl(): Promise<string | null> {
   const pendingImageReturn = await consumePendingChatImageReturn();
@@ -30,14 +31,28 @@ export async function peekSubscriptionExplicitExitUrl(): Promise<string | null>
 export function getSubscriptionFallbackExitUrl(
   returnTo: SubscriptionReturnTo,
 ): string {
-  return returnTo === "chat" ? ROUTES.chat : ROUTES.sidebar;
+  if (returnTo === "chat") return ROUTES.chat;
+  if (returnTo === "private-room") return ROUTES.privateRoom;
+  return ROUTES.sidebar;
 }
 
 export async function consumeSubscriptionExitUrl(
   returnTo: SubscriptionReturnTo,
 ): Promise<string> {
+  if (returnTo === "private-room") return ROUTES.privateRoom;
+
   return (
     (await consumeSubscriptionExplicitExitUrl()) ??
     getSubscriptionFallbackExitUrl(returnTo)
   );
 }
+
+export async function resolveSubscriptionSuccessExitUrl(
+  returnTo: SubscriptionReturnTo,
+): Promise<string> {
+  if (returnTo === "private-room") return ROUTES.privateRoom;
+
+  const pendingExitUrl = await peekSubscriptionExplicitExitUrl();
+  if (pendingExitUrl) return pendingExitUrl;
+  return consumeSubscriptionExitUrl(returnTo);
+}

src/lib/payment/__tests__/payment_search_params.test.ts

@@ -4,6 +4,7 @@ import {
   getFirstPaymentSearchParam,
   parsePaymentPayChannel,
   parsePaymentReturnSearchParams,
+  parseSubscriptionReturnTo,
 } from "../payment_search_params";
 
 describe("payment search params", () => {
@@ -35,4 +36,14 @@ describe("payment search params", () => {
       shouldResumePendingOrder: false,
     });
   });
+
+  it("accepts only supported subscription return targets", () => {
+    expect(parseSubscriptionReturnTo("chat")).toBe("chat");
+    expect(parseSubscriptionReturnTo("private-room")).toBe("private-room");
+    expect(parseSubscriptionReturnTo(["private-room", "chat"])).toBe(
+      "private-room",
+    );
+    expect(parseSubscriptionReturnTo("tip")).toBeNull();
+    expect(parseSubscriptionReturnTo(undefined)).toBeNull();
+  });
 });

src/lib/payment/__tests__/pending_payment_order.test.ts

@@ -1,6 +1,11 @@
 import { describe, expect, it } from "vitest";
 
-import { buildPendingPaymentSubscriptionUrl } from "../pending_payment_order";
+import {
+  buildPendingPaymentSubscriptionUrl,
+  clearPendingPaymentOrder,
+  getPendingPaymentOrder,
+  savePendingEzpayOrder,
+} from "../pending_payment_order";
 
 describe("pending payment order helpers", () => {
   it("keeps the pay channel when rebuilding subscription return urls", () => {
@@ -13,6 +18,40 @@ describe("pending payment order helpers", () => {
     ).toBe("/subscription?type=topup&payChannel=ezpay&paymentReturn=1&returnTo=chat");
   });
 
+  it("preserves private room returns when rebuilding Ezpay urls", () => {
+    expect(
+      buildPendingPaymentSubscriptionUrl({
+        payChannel: "ezpay",
+        returnTo: "private-room",
+        subscriptionType: "topup",
+      }),
+    ).toBe(
+      "/subscription?type=topup&payChannel=ezpay&paymentReturn=1&returnTo=private-room",
+    );
+  });
+
+  it("stores private room as an Ezpay return target", async () => {
+    await clearPendingPaymentOrder();
+
+    const saveResult = await savePendingEzpayOrder({
+      orderId: "order-private-room",
+      returnTo: "private-room",
+      subscriptionType: "topup",
+      createdAt: 1,
+    });
+    const storedResult = await getPendingPaymentOrder();
+
+    expect(saveResult.success).toBe(true);
+    expect(storedResult).toMatchObject({
+      success: true,
+      data: {
+        orderId: "order-private-room",
+        returnTo: "private-room",
+      },
+    });
+    await clearPendingPaymentOrder();
+  });
+
   it("routes tip payments back to the tip page", () => {
     expect(
       buildPendingPaymentSubscriptionUrl({

src/lib/payment/payment_search_params.ts

@@ -1,4 +1,5 @@
 import type { PayChannel } from "@/data/dto/payment";
+import type { AppSubscriptionReturnTo } from "@/router/navigation-types";
 
 export type PaymentSearchParamValue = string | string[] | undefined;
 export type PaymentSearchParams = Record<string, PaymentSearchParamValue>;
@@ -21,6 +22,14 @@ export function parsePaymentPayChannel(
   return channel === "ezpay" || channel === "stripe" ? channel : null;
 }
 
+export function parseSubscriptionReturnTo(
+  value: PaymentSearchParamValue,
+): AppSubscriptionReturnTo {
+  const returnTo = getFirstPaymentSearchParam(value);
+  if (returnTo === "chat" || returnTo === "private-room") return returnTo;
+  return null;
+}
+
 export function parsePaymentReturnSearchParams(
   searchParams: PaymentSearchParams,
 ): PaymentReturnSearchParams {

src/router/__tests__/navigation-resolver.test.ts

@@ -54,6 +54,17 @@ describe("navigation resolver", () => {
     ).toBe(ROUTE_BUILDERS.authWithRedirect(subscriptionUrl));
   });
 
+  it("builds private room subscription return urls", () => {
+    expect(
+      ROUTE_BUILDERS.subscription("topup", {
+        payChannel: "stripe",
+        returnTo: "private-room",
+      }),
+    ).toBe(
+      "/subscription?type=topup&payChannel=stripe&returnTo=private-room",
+    );
+  });
+
   it("allows not logged in users to enter chat for guest bootstrap", () => {
     expect(
       resolveRouteGuardRedirect({

src/router/navigation-types.ts

@@ -7,7 +7,7 @@ import type {
 } from "@/data/storage/navigation";
 
 export type AppSubscriptionType = "vip" | "topup";
-export type AppSubscriptionReturnTo = "chat" | null;
+export type AppSubscriptionReturnTo = "chat" | "private-room" | null;
 
 export interface OpenSubscriptionInput {
   type: AppSubscriptionType;

src/router/routes.ts

@@ -8,6 +8,7 @@
  */
 
 import type { PayChannel } from "@/data/dto/payment";
+import type { AppSubscriptionReturnTo } from "@/router/navigation-types";
 import {
   PAYMENT_ANALYTICS_ENTRY_PARAM,
   PAYMENT_ANALYTICS_REASON_PARAM,
@@ -36,7 +37,7 @@ export const ROUTE_BUILDERS = {
     type: "vip" | "topup",
     options: {
       payChannel?: PayChannel;
-      returnTo?: "chat";
+      returnTo?: Exclude<AppSubscriptionReturnTo, null>;
       analytics?: PaymentAnalyticsContext;
     } = {},
   ): `/subscription?${string}` => {

src/router/use-app-navigator.ts

@@ -7,7 +7,7 @@ import { NavigationStorage } from "@/data/storage/navigation";
 import type { SubscriptionReturnTo } from "@/lib/navigation/subscription_exit";
 import {
   consumeSubscriptionExitUrl,
-  peekSubscriptionExplicitExitUrl,
+  resolveSubscriptionSuccessExitUrl,
 } from "@/lib/navigation/subscription_exit";
 import { getDefaultPayChannelForCountryCode } from "@/lib/payment/default_pay_channel";
 import {
@@ -193,12 +193,7 @@ export function useAppNavigator(): AppNavigator {
 
   const exitSubscriptionAfterSuccess = useCallback((returnTo: SubscriptionReturnTo): void => {
     void (async () => {
-      const pendingExitUrl = await peekSubscriptionExplicitExitUrl();
-      if (pendingExitUrl) {
-        router.replace(pendingExitUrl);
-        return;
-      }
-      router.replace(await consumeSubscriptionExitUrl(returnTo));
+      router.replace(await resolveSubscriptionSuccessExitUrl(returnTo));
     })();
   }, [router]);