fix: restore puzzle runtime url state

src/components/platform-entry/usePlatformCreationAgentFlowController.ts

@@ -1,4 +1,5 @@
-import { useCallback, useRef, useState } from 'react';
+import { useCallback, useEffect, useRef, useState } from 'react';
+import type { Dispatch, SetStateAction } from 'react';
 
 import type { TextStreamOptions } from '../../services/aiTypes';
 import type { SelectionStage } from './platformEntryTypes';
@@ -75,12 +76,13 @@ type PlatformCreationAgentFlowControllerOptions<
   enterCreateTab: () => void;
   setSelectionStage: (stage: SelectionStage) => void;
   onSessionOpened?: () => void;
+  onSessionChanged?: (session: TSession | null) => void;
   onOpenError?: (params: { error: unknown; errorMessage: string }) => void;
   onActionComplete?: (params: {
     payload: TActionPayload;
     response: TActionResponse;
     session: TSession;
-    setSession: (session: TSession) => void;
+    setSession: Dispatch<SetStateAction<TSession | null>>;
   }) =>
     | Promise<{ openResult?: boolean } | void>
     | { openResult?: boolean }
@@ -94,7 +96,7 @@ type PlatformCreationAgentFlowControllerOptions<
     error: unknown;
     errorMessage: string;
     session: TSession;
-    setSession: (session: TSession) => void;
+    setSession: Dispatch<SetStateAction<TSession | null>>;
   }) => void | Promise<void>;
 };
 
@@ -141,12 +143,27 @@ export function usePlatformCreationAgentFlowController<
     TActionResponse
   >,
 ) {
-  const [session, setSession] = useState<TSession | null>(null);
+  const [session, rawSetSession] = useState<TSession | null>(null);
   const [error, setError] = useState<string | null>(null);
   const [isBusy, setIsBusy] = useState(false);
   const [streamingReplyText, setStreamingReplyText] = useState('');
   const [isStreamingReply, setIsStreamingReply] = useState(false);
   const latestStreamingReplyTextRef = useRef('');
+  const onSessionChangedRef = useRef(options.onSessionChanged);
+
+  useEffect(() => {
+    onSessionChangedRef.current = options.onSessionChanged;
+  }, [options.onSessionChanged]);
+
+  const setSession = useCallback(
+    (nextSessionOrUpdater: SetStateAction<TSession | null>) => {
+      rawSetSession(nextSessionOrUpdater);
+      if (typeof nextSessionOrUpdater !== 'function') {
+        onSessionChangedRef.current?.(nextSessionOrUpdater);
+      }
+    },
+    [],
+  );
 
   const updateStreamingReplyText = useCallback((text: string) => {
     latestStreamingReplyTextRef.current = text;
@@ -174,10 +191,10 @@ export function usePlatformCreationAgentFlowController<
           createPayload ?? options.createPayload,
         );
         const nextSession = options.client.selectSession(response);
-        setSession(nextSession);
         options.enterCreateTab();
         options.onSessionOpened?.();
         options.setSelectionStage(options.workspaceStage);
+        setSession(nextSession);
         return nextSession;
       } catch (caughtError) {
         const errorMessage = options.resolveErrorMessage(
@@ -212,11 +229,11 @@ export function usePlatformCreationAgentFlowController<
       try {
         const response = await options.client.getSession(normalizedSessionId);
         const nextSession = options.client.selectSession(response);
-        setSession(nextSession);
         options.enterCreateTab();
         options.setSelectionStage(
           nextSession.draft ? options.resultStage : options.workspaceStage,
         );
+        setSession(nextSession);
         return nextSession;
       } catch (caughtError) {
         setError(