fix(app): fork should copy prompt into new session (#17375)

3 files changed, 14 insertions, 10 deletions

diff --git a/packages/app/src/components/dialog-fork.tsx b/packages/app/src/components/dialog-fork.tsx
index 8810955cc..9e1b896fa 100644
--- a/packages/app/src/components/dialog-fork.tsx
+++ b/packages/app/src/components/dialog-fork.tsx
@@ -66,6 +66,7 @@ export const DialogFork: Component = () => {
       directory: sdk.directory,
       attachmentName: language.t("common.attachment"),
     })
+    const dir = base64Encode(sdk.directory)
 
     sdk.client.session
       .fork({ sessionID, messageID: item.id })
@@ -75,10 +76,8 @@ export const DialogFork: Component = () => {
           return
         }
         dialog.close()
-        navigate(`/${base64Encode(sdk.directory)}/session/${forked.data.id}`)
-        requestAnimationFrame(() => {
-          prompt.set(restored)
-        })
+        prompt.set(restored, undefined, { dir, id: forked.data.id })
+        navigate(`/${dir}/session/${forked.data.id}`)
       })
       .catch((err: unknown) => {
         const message = err instanceof Error ? err.message : String(err)
diff --git a/packages/app/src/context/prompt.tsx b/packages/app/src/context/prompt.tsx
index fb8226559..4aeb05e61 100644
--- a/packages/app/src/context/prompt.tsx
+++ b/packages/app/src/context/prompt.tsx
@@ -151,6 +151,11 @@ const MAX_PROMPT_SESSIONS = 20
 
 type PromptSession = ReturnType<typeof createPromptSession>
 
+type Scope = {
+  dir: string
+  id?: string
+}
+
 type PromptCacheEntry = {
   value: PromptSession
   dispose: VoidFunction
@@ -265,6 +270,7 @@ export const { use: usePrompt, provider: PromptProvider } = createSimpleContext(
     }
 
     const session = createMemo(() => load(params.dir!, params.id))
+    const pick = (scope?: Scope) => (scope ? load(scope.dir, scope.id) : session())
 
     return {
       ready: () => session().ready(),
@@ -280,8 +286,8 @@ export const { use: usePrompt, provider: PromptProvider } = createSimpleContext(
           session().context.updateComment(path, commentID, next),
         replaceComments: (items: FileContextItem[]) => session().context.replaceComments(items),
       },
-      set: (prompt: Prompt, cursorPosition?: number) => session().set(prompt, cursorPosition),
-      reset: () => session().reset(),
+      set: (prompt: Prompt, cursorPosition?: number, scope?: Scope) => pick(scope).set(prompt, cursorPosition),
+      reset: (scope?: Scope) => pick(scope).reset(),
     }
   },
 })
diff --git a/packages/app/src/pages/session.tsx b/packages/app/src/pages/session.tsx
index d917ce4c7..8399a1367 100644
--- a/packages/app/src/pages/session.tsx
+++ b/packages/app/src/pages/session.tsx
@@ -1477,6 +1477,7 @@ export default function Page() {
 
   const fork = (input: { sessionID: string; messageID: string }) => {
     const value = draft(input.messageID)
+    const dir = base64Encode(sdk.directory)
     return sdk.client.session
       .fork(input)
       .then((result) => {
@@ -1488,10 +1489,8 @@ export default function Page() {
           })
           return
         }
-        navigate(`/${base64Encode(sdk.directory)}/session/${next.id}`)
-        requestAnimationFrame(() => {
-          prompt.set(value)
-        })
+        prompt.set(value, undefined, { dir, id: next.id })
+        navigate(`/${dir}/session/${next.id}`)
       })
       .catch(fail)
   }