From 9adb0fee28a89d7c8618f6175add80120d8ba7b0 Mon Sep 17 00:00:00 2001
From: xuhd <hdx@kland.com.cn>
Date: Fri, 5 Jun 2026 13:58:34 +0800
Subject: [PATCH] fix(notebook): preserve selected notebook/docs when switching
 material source with studioToken

- Initialize materialSource via WritingFormProvider initialMaterialSource prop
- Don't clear selectedNotebookId/selectedDocumentIds on source switch when studioToken is set

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 .../components/ai-writing/AIWritingForm.tsx     | 17 +++++++----------
 .../open-canvas/contexts/WritingFormContext.tsx |  4 ++--
 .../src/open-canvas/pages/AIWritingPage.tsx     |  2 +-
 3 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/frontend-web/src/open-canvas/components/ai-writing/AIWritingForm.tsx b/frontend-web/src/open-canvas/components/ai-writing/AIWritingForm.tsx
index 91bee967..9849b9ff 100644
--- a/frontend-web/src/open-canvas/components/ai-writing/AIWritingForm.tsx
+++ b/frontend-web/src/open-canvas/components/ai-writing/AIWritingForm.tsx
@@ -97,10 +97,6 @@ export function AIWritingForm({
     [models, modelName],
   );
 
-  // studioToken / initialNotebookId 存在时强制使用「我的笔记」
-  const effectiveMaterialSource: MaterialSource = studioToken ?? initialNotebookId != null
-    ? 'notebook'
-    : materialSource;
 
   // 初始化时加载文章类型
   useEffect(() => {
@@ -143,7 +139,8 @@ export function AIWritingForm({
 
   function handleMaterialSourceChange(val: MaterialSource) {
     setMaterialSource(val);
-    if (val !== 'notebook') {
+    // studioToken 模式下笔记本和文档已锁定，切换素材来源时不清空
+    if (val !== 'notebook' && !studioToken) {
       setSelectedNotebookId(null);
       setSelectedDocumentIds([]);
     }
@@ -169,11 +166,11 @@ export function AIWritingForm({
       wordCountTarget: wordCount,
       keywordCount: DEFAULT_KEYWORD_COUNT,
       writingMode: hideAdvancedFields ? 'loose' : strictMode ? 'strict' : 'loose',
-      materialSource: effectiveMaterialSource,
+      materialSource: materialSource,
       modelName: selectedModel?.name || modelName || undefined,
       // 只有选了"我的笔记"才传笔记参数
-      noteSearchSpaceId: effectiveMaterialSource === 'notebook' ? (selectedNotebookId ?? undefined) : undefined,
-      noteSourceIds: effectiveMaterialSource === 'notebook' && selectedNotebookId != null ? selectedDocumentIds : undefined,
+      noteSearchSpaceId: materialSource === 'notebook' ? (selectedNotebookId ?? undefined) : undefined,
+      noteSourceIds: materialSource === 'notebook' && selectedNotebookId != null ? selectedDocumentIds : undefined,
     });
   }
 
@@ -282,7 +279,7 @@ export function AIWritingForm({
         <div className="mt-4 space-y-2">
           <label className={FIELD_LABEL_CLASS}>素材来源</label>
           <Select
-            value={effectiveMaterialSource}
+            value={materialSource}
             onValueChange={(v) => handleMaterialSourceChange(v as MaterialSource)}
             disabled={disabled}
           >
@@ -300,7 +297,7 @@ export function AIWritingForm({
           </Select>
 
           {/* 选了"我的笔记"才展示笔记本 / 文档选择器 */}
-          {effectiveMaterialSource === 'notebook' && (
+          {materialSource === 'notebook' && (
             <div className="mt-2 space-y-2">
               {/* studioToken 模式：笔记本已锁定，不展示选择器 */}
               {!studioToken && (
diff --git a/frontend-web/src/open-canvas/contexts/WritingFormContext.tsx b/frontend-web/src/open-canvas/contexts/WritingFormContext.tsx
index 52bcb331..41fe9198 100644
--- a/frontend-web/src/open-canvas/contexts/WritingFormContext.tsx
+++ b/frontend-web/src/open-canvas/contexts/WritingFormContext.tsx
@@ -53,7 +53,7 @@ export function useWritingFormOptional(): WritingFormValue | null {
   return useContext(WritingFormContext);
 }
 
-export function WritingFormProvider({ children }: { children: ReactNode }) {
+export function WritingFormProvider({ children, initialMaterialSource }: { children: ReactNode; initialMaterialSource?: 'knowledge_base' | 'skill' | 'notebook' }) {
   const [intent, setIntent] = useState('');
   const [showOutline, setShowOutline] = useState(false);
   const [outlineText, setOutlineText] = useState('');
@@ -61,7 +61,7 @@ export function WritingFormProvider({ children }: { children: ReactNode }) {
   const [wordCountInput, setWordCountInput] = useState('800');
   const [audience, setAudience] = useState('普通读者');
   const [strictMode, setStrictMode] = useState(false);
-  const [materialSource, setMaterialSource] = useState<'knowledge_base' | 'skill' | 'notebook'>('knowledge_base');
+  const [materialSource, setMaterialSource] = useState<'knowledge_base' | 'skill' | 'notebook'>(initialMaterialSource ?? 'knowledge_base');
   const [modelName, setModelName] = useState('');
 
   const applyProposal = useCallback((payload: SetupWritingApprovalPayload) => {
diff --git a/frontend-web/src/open-canvas/pages/AIWritingPage.tsx b/frontend-web/src/open-canvas/pages/AIWritingPage.tsx
index 5d32da94..34c323db 100644
--- a/frontend-web/src/open-canvas/pages/AIWritingPage.tsx
+++ b/frontend-web/src/open-canvas/pages/AIWritingPage.tsx
@@ -223,7 +223,7 @@ export default function AIWritingPage({
       <CanvasUIProvider>
       <AIWritingHistoryProvider>
       <AIWritingProvider>
-      <WritingFormProvider>
+      <WritingFormProvider initialMaterialSource={studioToken ?? initialNotebookId != null ? 'notebook' : undefined}>
       <TooltipProvider delayDuration={300}>
         <div className="h-full flex flex-col overflow-hidden bg-background">
           <div className="flex-1 min-h-0 overflow-hidden">