From 028568c3ae9ce2387ece9525835df8b3b942b2cd Mon Sep 17 00:00:00 2001
From: Supra4E8C <tan13318527080@163.com>
Date: Sun, 22 Feb 2026 02:51:09 +0800
Subject: [PATCH] fix(editing): harden unsaved navigation and draft cleanup

---
 src/hooks/useUnsavedChangesGuard.ts    | 14 ++++++++------
 src/pages/AiProvidersCodexEditPage.tsx |  3 ++-
 src/stores/useClaudeEditDraftStore.ts  | 10 ++++++----
 src/stores/useOpenAIEditDraftStore.ts  | 10 ++++++----
 4 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/src/hooks/useUnsavedChangesGuard.ts b/src/hooks/useUnsavedChangesGuard.ts
index 14e2cf2..8a97508 100644
--- a/src/hooks/useUnsavedChangesGuard.ts
+++ b/src/hooks/useUnsavedChangesGuard.ts
@@ -23,19 +23,21 @@ export function useUnsavedChangesGuard(options: UseUnsavedChangesGuardOptions) {
   const { enabled = true, shouldBlock, dialog } = options;
   const { showConfirmation } = useNotificationStore();
   const lastBlockedRef = useRef<string>('');
-  const allowNextNavigationRef = useRef(false);
+  const allowNextNavigationUntilRef = useRef(0);
   const location = useLocation();
 
   const allowNextNavigation = useCallback(() => {
-    allowNextNavigationRef.current = true;
+    // Allow a short window for programmatic navigations after successful save.
+    // This avoids stale "allow" flags lingering when no navigation happens.
+    allowNextNavigationUntilRef.current = Date.now() + 2_000;
   }, []);
 
   const shouldBlockFunction = useCallback<BlockerFunction>(
     (args) => {
-      if (allowNextNavigationRef.current) {
+      if (!enabled) return false;
+      if (allowNextNavigationUntilRef.current > Date.now()) {
         return false;
       }
-      if (!enabled) return false;
       return typeof shouldBlock === 'function' ? shouldBlock(args) : shouldBlock;
     },
     [enabled, shouldBlock]
@@ -44,8 +46,8 @@ export function useUnsavedChangesGuard(options: UseUnsavedChangesGuardOptions) {
   const blocker = useBlocker(shouldBlockFunction);
 
   useEffect(() => {
-    if (!allowNextNavigationRef.current) return;
-    allowNextNavigationRef.current = false;
+    if (allowNextNavigationUntilRef.current === 0) return;
+    allowNextNavigationUntilRef.current = 0;
   }, [location.key]);
 
   const blockedKey = useMemo(() => {
diff --git a/src/pages/AiProvidersCodexEditPage.tsx b/src/pages/AiProvidersCodexEditPage.tsx
index 6691fd2..85bd4a5 100644
--- a/src/pages/AiProvidersCodexEditPage.tsx
+++ b/src/pages/AiProvidersCodexEditPage.tsx
@@ -185,6 +185,7 @@ export function AiProvidersCodexEditPage() {
     if (initialData) {
       const nextForm: ProviderFormState = {
         ...initialData,
+        websockets: Boolean(initialData.websockets),
         headers: headersToEntries(initialData.headers),
         modelEntries: modelsToEntries(initialData.models),
         excludedText: excludedModelsToText(initialData.excludedModels),
@@ -362,7 +363,7 @@ export function AiProvidersCodexEditPage() {
         priority: form.priority !== undefined ? Math.trunc(form.priority) : undefined,
         prefix: form.prefix?.trim() || undefined,
         baseUrl,
-        websockets: form.websockets ?? false,
+        websockets: Boolean(form.websockets),
         proxyUrl: form.proxyUrl?.trim() || undefined,
         headers: buildHeaderObject(form.headers),
         models: entriesToModels(form.modelEntries),
diff --git a/src/stores/useClaudeEditDraftStore.ts b/src/stores/useClaudeEditDraftStore.ts
index d7ab7aa..f0bd7f2 100644
--- a/src/stores/useClaudeEditDraftStore.ts
+++ b/src/stores/useClaudeEditDraftStore.ts
@@ -196,10 +196,12 @@ export const useClaudeEditDraftStore = create<ClaudeEditDraftState>((set, get) =
   clearDraft: (key) => {
     if (!key) return;
     set((state) => {
-      if (!state.drafts[key]) return state;
-      const next = { ...state.drafts };
-      delete next[key];
-      return { drafts: next };
+      if (!state.drafts[key] && !state.refCounts[key]) return state;
+      const nextDrafts = { ...state.drafts };
+      delete nextDrafts[key];
+      const nextCounts = { ...state.refCounts };
+      delete nextCounts[key];
+      return { drafts: nextDrafts, refCounts: nextCounts };
     });
   },
 }));
diff --git a/src/stores/useOpenAIEditDraftStore.ts b/src/stores/useOpenAIEditDraftStore.ts
index cff52a0..fa4fb23 100644
--- a/src/stores/useOpenAIEditDraftStore.ts
+++ b/src/stores/useOpenAIEditDraftStore.ts
@@ -227,10 +227,12 @@ export const useOpenAIEditDraftStore = create<OpenAIEditDraftState>((set, get) =
   clearDraft: (key) => {
     if (!key) return;
     set((state) => {
-      if (!state.drafts[key]) return state;
-      const next = { ...state.drafts };
-      delete next[key];
-      return { drafts: next };
+      if (!state.drafts[key] && !state.refCounts[key]) return state;
+      const nextDrafts = { ...state.drafts };
+      delete nextDrafts[key];
+      const nextCounts = { ...state.refCounts };
+      delete nextCounts[key];
+      return { drafts: nextDrafts, refCounts: nextCounts };
     });
   },
 }));