diff --git a/src/app-server/client.ts b/src/app-server/client.ts index 200b43ae..d048b03d 100644 --- a/src/app-server/client.ts +++ b/src/app-server/client.ts @@ -257,8 +257,8 @@ export class AppServerClient { return this.request("thread/unarchive", { threadId }); } - rollbackThread(threadId: string): Promise { - return this.request("thread/rollback", { threadId, numTurns: 1 }); + rollbackThread(threadId: string, numTurns = 1): Promise { + return this.request("thread/rollback", { threadId, numTurns }); } setThreadName(threadId: string, name: string): Promise { diff --git a/src/features/chat/chat-message-renderer.ts b/src/features/chat/chat-message-renderer.ts index cac5f385..aa01f760 100644 --- a/src/features/chat/chat-message-renderer.ts +++ b/src/features/chat/chat-message-renderer.ts @@ -8,6 +8,7 @@ import type { ComposerBoundaryScrollAction } from "./composer/boundary-scroll"; import { MessageScrollController, type MessageScrollIntent } from "./ui/scroll"; import type { ChatTurnDiffViewState } from "./ui/turn-diff"; import { MarkdownMessageRenderer } from "./markdown-message-renderer"; +import { forkCandidatesFromItems, isForkCandidateItem } from "./fork"; import { isRollbackCandidateItem, rollbackCandidateFromItems } from "./rollback"; import { chatTurnBusy, type ChatAction, type ChatState, type ChatStateStore } from "./chat-state"; import { implementPlanCandidateFromState } from "./plan-implementation"; @@ -21,6 +22,7 @@ export interface ChatMessageRendererOptions { consumeScrollIntent: () => MessageScrollIntent; loadOlderTurns: () => void; rollbackThread: (threadId: string) => void; + forkThreadFromTurn: (threadId: string, turnId: string, archiveSource: boolean) => void; implementPlan: (item: DisplayItem) => void; openTurnDiff: (state: ChatTurnDiffViewState) => void; pendingRequestsSignature: () => string; @@ -64,6 +66,7 @@ export class ChatMessageRenderer { const scrollPlan = this.scrollController.prepareRender(messagesEl, this.options.consumeScrollIntent()); const busy = chatTurnBusy(state); const rollbackCandidate = busy ? null : rollbackCandidateFromItems(state.displayItems); + const forkCandidates = busy ? [] : forkCandidatesFromItems(state.displayItems); const implementPlanCandidate = implementPlanCandidateFromState(state); const blocks = messageStreamBlocks({ @@ -96,6 +99,10 @@ export class ChatMessageRenderer { onRollbackItem: () => { if (state.activeThreadId) this.options.rollbackThread(state.activeThreadId); }, + canForkItem: (item: DisplayItem) => isForkCandidateItem(item, forkCandidates), + onForkItem: (item, archiveSource) => { + if (state.activeThreadId && item.turnId) this.options.forkThreadFromTurn(state.activeThreadId, item.turnId, archiveSource); + }, openTurnDiff: (turnDiffState) => { this.options.openTurnDiff(turnDiffState); }, @@ -131,6 +138,13 @@ export class ChatMessageRenderer { } private setOpenDetail(key: string, open: boolean): void { + if (open && key.startsWith("message:fork-actions:")) { + for (const openKey of this.state.openDetails) { + if (openKey.startsWith("message:fork-actions:") && openKey !== key) { + this.dispatch({ type: "ui/detail-open-set", key: openKey, open: false }); + } + } + } this.dispatch({ type: "ui/detail-open-set", key, open }); } } diff --git a/src/features/chat/chat-view-controller-assembly.ts b/src/features/chat/chat-view-controller-assembly.ts index 28828303..8f91af29 100644 --- a/src/features/chat/chat-view-controller-assembly.ts +++ b/src/features/chat/chat-view-controller-assembly.ts @@ -108,6 +108,7 @@ export interface ChatViewControllerAssemblyHost { registerActiveLeafChange: (handler: (leaf: WorkspaceLeaf | null) => void) => void; isOwnLeaf: (leaf: WorkspaceLeaf | null) => boolean; archiveAdapter: () => ArchiveExportAdapter; + closePanel: () => void; } export function createChatViewControllerAssembly(host: ChatViewControllerAssemblyHost): ChatViewControllerAssembly { @@ -200,6 +201,7 @@ export function createChatViewControllerAssembly(host: ChatViewControllerAssembl consumeScrollIntent: () => host.messageScroll.consumeIntent(), loadOlderTurns: () => void history.loadOlder(), rollbackThread: (threadId) => void threadActions.rollbackThread(threadId), + forkThreadFromTurn: (threadId, turnId, archiveSource) => void threadActions.forkThreadFromTurn(threadId, turnId, archiveSource), implementPlan: (item) => void planImplementation.implement(item), openTurnDiff: (state) => void host.plugin.openTurnDiff(state), pendingRequestsSignature: host.pendingRequestsSignature, @@ -397,6 +399,7 @@ export function createChatViewControllerAssembly(host: ChatViewControllerAssembl refreshSharedThreadListFromOpenSurface: () => { host.plugin.refreshSharedThreadListFromOpenSurface(); }, + closePanel: host.closePanel, }); toolbarPanels = new ToolbarPanelController({ stateStore: host.stateStore, diff --git a/src/features/chat/display/blocks.ts b/src/features/chat/display/blocks.ts index 163ad335..96a1440c 100644 --- a/src/features/chat/display/blocks.ts +++ b/src/features/chat/display/blocks.ts @@ -1,4 +1,5 @@ import type { DisplayBlock, DisplayItem, DisplayKind } from "./types"; +import { isFinalAssistantMessage } from "./final-assistant"; import { pathRelativeToRoot } from "./paths"; import { executionState } from "./state"; @@ -73,10 +74,6 @@ function finalAssistantItemsByTurn(items: readonly DisplayItem[]): Map, diff --git a/src/features/chat/display/final-assistant.ts b/src/features/chat/display/final-assistant.ts new file mode 100644 index 00000000..5c21e729 --- /dev/null +++ b/src/features/chat/display/final-assistant.ts @@ -0,0 +1,5 @@ +import type { DisplayItem } from "./types"; + +export function isFinalAssistantMessage(item: DisplayItem): boolean { + return item.kind === "message" && item.role === "assistant" && item.markdown !== false; +} diff --git a/src/features/chat/fork.ts b/src/features/chat/fork.ts new file mode 100644 index 00000000..367c1de1 --- /dev/null +++ b/src/features/chat/fork.ts @@ -0,0 +1,37 @@ +import type { DisplayItem } from "./display/types"; +import { isFinalAssistantMessage } from "./display/final-assistant"; + +export interface ForkCandidate { + itemId: string; + turnId: string; +} + +export function forkCandidatesFromItems(items: readonly DisplayItem[]): readonly ForkCandidate[] { + const finalAssistantItemsByTurn = new Map(); + for (const item of items) { + if (!item.turnId || !isFinalAssistantMessage(item)) continue; + finalAssistantItemsByTurn.set(item.turnId, { itemId: item.id, turnId: item.turnId }); + } + return [...finalAssistantItemsByTurn.values()]; +} + +export function isForkCandidateItem(item: DisplayItem, candidates: readonly ForkCandidate[]): boolean { + return candidates.some((candidate) => item.id === candidate.itemId && item.turnId === candidate.turnId); +} + +export function turnsAfterTurnId(items: readonly DisplayItem[], turnId: string): number | null { + const turnIds = orderedTurnIds(items); + const index = turnIds.indexOf(turnId); + return index === -1 ? null : turnIds.length - index - 1; +} + +function orderedTurnIds(items: readonly DisplayItem[]): string[] { + const turnIds: string[] = []; + const seen = new Set(); + for (const item of items) { + if (!item.turnId || seen.has(item.turnId)) continue; + seen.add(item.turnId); + turnIds.push(item.turnId); + } + return turnIds; +} diff --git a/src/features/chat/thread-actions.ts b/src/features/chat/thread-actions.ts index da4b5978..932b59ba 100644 --- a/src/features/chat/thread-actions.ts +++ b/src/features/chat/thread-actions.ts @@ -7,6 +7,7 @@ import type { CodexPanelSettings } from "../../settings/model"; import type { ArchiveExportAdapter } from "../../domain/threads/export"; import { chatTurnBusy, type ChatAction, type ChatState, type ChatStateStore } from "./chat-state"; import type { ThreadHistoryLoader } from "./thread-history"; +import { turnsAfterTurnId } from "./fork"; import { rollbackCandidateFromItems } from "./rollback"; export interface ChatThreadActionControllerHost { @@ -26,6 +27,7 @@ export interface ChatThreadActionControllerHost { notifyActiveThreadIdentityChanged: () => void; refreshThreads: () => Promise; refreshSharedThreadListFromOpenSurface: () => void; + closePanel: () => void; } export class ChatThreadActionController { @@ -40,12 +42,16 @@ export class ChatThreadActionController { } async archiveThread(threadId: string, saveMarkdown = this.host.settings().archiveExportEnabled): Promise { + await this.archiveThreadWithResult(threadId, saveMarkdown); + } + + private async archiveThreadWithResult(threadId: string, saveMarkdown = this.host.settings().archiveExportEnabled): Promise { if (chatTurnBusy(this.state)) { this.host.addSystemMessage("Finish or interrupt the current turn before archiving threads."); - return; + return false; } const client = this.host.currentClient(); - if (!client) return; + if (!client) return false; try { const settings = this.host.settings(); if (saveMarkdown) { @@ -59,12 +65,18 @@ export class ChatThreadActionController { } await client.archiveThread(threadId); this.host.notifyThreadArchived(threadId); + return true; } catch (error) { this.host.addSystemMessage(error instanceof Error ? error.message : String(error)); + return false; } } async forkThread(threadId: string): Promise { + await this.forkThreadFromTurn(threadId, null, false); + } + + async forkThreadFromTurn(threadId: string, turnId: string | null, archiveSource: boolean): Promise { if (chatTurnBusy(this.state)) { this.host.addSystemMessage("Finish or interrupt the current turn before forking threads."); return; @@ -73,10 +85,19 @@ export class ChatThreadActionController { const client = this.host.currentClient(); if (!client) return; + const turnsToDrop = turnId ? turnsAfterTurnId(this.state.displayItems, turnId) : 0; + if (turnsToDrop === null) { + this.host.addSystemMessage("Could not find the selected turn to fork."); + return; + } + try { const sourceName = inheritedForkThreadName(threadId, this.state.listedThreads); const response = await client.forkThread(threadId, this.host.vaultPath); const forkedThreadId = response.thread.id; + if (turnsToDrop > 0) { + await client.rollbackThread(forkedThreadId, turnsToDrop); + } if (sourceName) { try { await client.setThreadName(forkedThreadId, sourceName); @@ -86,12 +107,19 @@ export class ChatThreadActionController { this.host.addSystemMessage(`Forked thread ${forkedThreadId}, but could not copy the source thread name: ${message}`); } } + let openedForkPanel = false; try { await this.host.openThreadInNewView(forkedThreadId); + openedForkPanel = true; } catch (error) { const message = error instanceof Error ? error.message : String(error); this.host.addSystemMessage(`Forked thread ${forkedThreadId}, but could not open it in a new panel: ${message}`); } + if (archiveSource) { + if (!openedForkPanel) return; + const archived = await this.archiveThreadWithResult(threadId); + if (archived) this.host.closePanel(); + } } catch (error) { this.host.addSystemMessage(error instanceof Error ? error.message : String(error)); } diff --git a/src/features/chat/ui/message-stream.tsx b/src/features/chat/ui/message-stream.tsx index fafc714f..aaafd66c 100644 --- a/src/features/chat/ui/message-stream.tsx +++ b/src/features/chat/ui/message-stream.tsx @@ -1,4 +1,4 @@ -import { Fragment, forwardRef, useLayoutEffect, useRef, useState, type ReactNode } from "react"; +import { Fragment, forwardRef, useEffect, useLayoutEffect, useRef, useState, type ReactNode } from "react"; import { displayBlocksForItems } from "../display/blocks"; import { executionState } from "../display/state"; @@ -37,6 +37,8 @@ export interface MessageStreamContext { onImplementPlanItem?: (item: DisplayItem) => void; canRollbackItem?: (item: DisplayItem) => boolean; onRollbackItem?: (item: DisplayItem) => void; + canForkItem?: (item: DisplayItem) => boolean; + onForkItem?: (item: DisplayItem, archiveSource: boolean) => void; openTurnDiff?: (state: ChatTurnDiffViewState) => void; pendingRequestsSignature?: string; renderPendingRequests?: () => ReactNode; @@ -221,17 +223,50 @@ function MessageItem({ item, context }: { item: RenderableMessageItem; context: } function MessageRole({ item, context }: { item: RenderableMessageItem; context: MessageStreamContext }): ReactNode { + const forkActionsKey = `message:fork-actions:${item.id}`; + const forkActionsOpen = context.openDetails.has(forkActionsKey); + const roleRef = useRef(null); + + useEffect(() => { + if (!forkActionsOpen) return; + const doc = roleRef.current?.ownerDocument; + if (!doc) return; + const closeOnOutsidePointer = (event: PointerEvent) => { + if (event.target instanceof Node && roleRef.current?.contains(event.target)) return; + context.onDetailsToggle?.(forkActionsKey, false); + }; + doc.addEventListener("pointerdown", closeOnOutsidePointer, true); + return () => { + doc.removeEventListener("pointerdown", closeOnOutsidePointer, true); + }; + }, [context, forkActionsKey, forkActionsOpen]); + + const copyAction = + item.kind === "message" && context.copyText && isMessageCopyActionVisible(item, context) && !forkActionsOpen ? ( + context.copyText?.(item.copyText ?? item.text)} + /> + ) : null; + return ( -
+
{displayRoleLabel(item)} - {item.kind === "message" && context.copyText && isMessageCopyActionVisible(item, context) ? ( + {forkActionsOpen && context.canForkItem?.(item) ? ( context.copyText?.(item.copyText ?? item.text)} + icon="archive" + label="Fork and archive" + className="codex-panel__fork-and-archive-message" + onClick={() => { + context.onDetailsToggle?.(forkActionsKey, false); + context.onForkItem?.(item, true); + }} /> - ) : null} + ) : ( + copyAction + )} {context.canImplementPlanItem?.(item) ? ( context.onRollbackItem?.(item)} /> ) : null} + {context.canForkItem?.(item) ? ( + { + if (forkActionsOpen) { + context.onDetailsToggle?.(forkActionsKey, false); + context.onForkItem?.(item, false); + } else { + context.onDetailsToggle?.(forkActionsKey, true); + } + }} + /> + ) : null}
); } diff --git a/src/features/chat/view.ts b/src/features/chat/view.ts index aea1872f..5cc4503b 100644 --- a/src/features/chat/view.ts +++ b/src/features/chat/view.ts @@ -161,6 +161,9 @@ export class CodexChatView extends ItemView { }, isOwnLeaf: (candidateLeaf) => candidateLeaf === this.leaf, archiveAdapter: () => this.app.vault.adapter, + closePanel: () => { + this.leaf.detach(); + }, }); this.connection = controllers.connection; this.controller = controllers.controller; diff --git a/styles.css b/styles.css index fd6d9f57..aa43213f 100644 --- a/styles.css +++ b/styles.css @@ -1393,6 +1393,7 @@ } .codex-panel__message:hover .codex-panel__message-action, +.codex-panel__message-role--fork-open .codex-panel__message-action, .codex-panel__message-action:focus-visible { opacity: 1; } diff --git a/tests/features/chat/chat-message-renderer.test.ts b/tests/features/chat/chat-message-renderer.test.ts index a91a8a63..3c1351b7 100644 --- a/tests/features/chat/chat-message-renderer.test.ts +++ b/tests/features/chat/chat-message-renderer.test.ts @@ -243,6 +243,7 @@ function chatMessageRenderer( consumeScrollIntent: () => "auto", loadOlderTurns: vi.fn(), rollbackThread: vi.fn(), + forkThreadFromTurn: vi.fn(), implementPlan: vi.fn(), openTurnDiff: vi.fn(), pendingRequestsSignature: () => "", diff --git a/tests/features/chat/fork.test.ts b/tests/features/chat/fork.test.ts new file mode 100644 index 00000000..4ba9034d --- /dev/null +++ b/tests/features/chat/fork.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from "vitest"; + +import { forkCandidatesFromItems, isForkCandidateItem, turnsAfterTurnId } from "../../../src/features/chat/fork"; +import type { DisplayItem } from "../../../src/features/chat/display/types"; + +describe("fork candidates", () => { + it("selects final assistant messages and counts later turns", () => { + const items: DisplayItem[] = [ + { id: "u1", kind: "message", role: "user", text: "first", turnId: "turn-1", markdown: true }, + { id: "a1", kind: "message", role: "assistant", text: "first answer", turnId: "turn-1", markdown: true }, + { id: "tool-1", kind: "tool", role: "tool", text: "work", turnId: "turn-2" }, + { id: "a2-delta", kind: "message", role: "assistant", text: "draft", turnId: "turn-2", markdown: false }, + { id: "a2", kind: "message", role: "assistant", text: "second answer", turnId: "turn-2", markdown: true }, + { id: "u3", kind: "message", role: "user", text: "third", turnId: "turn-3", markdown: true }, + { id: "a3", kind: "message", role: "assistant", text: "third answer", turnId: "turn-3", markdown: true }, + ]; + + const candidates = forkCandidatesFromItems(items); + + expect(candidates).toEqual([ + { itemId: "a1", turnId: "turn-1" }, + { itemId: "a2", turnId: "turn-2" }, + { itemId: "a3", turnId: "turn-3" }, + ]); + expect(isForkCandidateItem(expectPresent(items[4]), candidates)).toBe(true); + expect(isForkCandidateItem(expectPresent(items[3]), candidates)).toBe(false); + expect(turnsAfterTurnId(items, "turn-1")).toBe(2); + expect(turnsAfterTurnId(items, "turn-2")).toBe(1); + expect(turnsAfterTurnId(items, "turn-3")).toBe(0); + expect(turnsAfterTurnId(items, "missing")).toBeNull(); + }); +}); + +function expectPresent(value: T | null | undefined): T { + if (value === null || value === undefined) throw new Error("Expected value to be present"); + return value; +} diff --git a/tests/features/chat/ui/message-stream.test.tsx b/tests/features/chat/ui/message-stream.test.tsx index 5c6f50b8..6ad2bc1f 100644 --- a/tests/features/chat/ui/message-stream.test.tsx +++ b/tests/features/chat/ui/message-stream.test.tsx @@ -518,6 +518,102 @@ describe("message stream block identity and message actions", () => { expect(copyText).toHaveBeenNthCalledWith(2, "# Answer"); }); + it("expands assistant fork actions in the copy slot and defaults repeat clicks to plain fork", () => { + const onDetailsToggle = vi.fn(); + const onForkItem = vi.fn(); + const item: DisplayItem = { + id: "a1", + kind: "message", + role: "assistant", + text: "answer", + copyText: "answer", + turnId: "turn-1", + markdown: true, + }; + + const closedBlock = messageStreamBlocks({ + activeThreadId: "thread", + turnLifecycle: idleTurnLifecycle(), + historyCursor: null, + loadingHistory: false, + displayItems: [item], + openDetails: new Set(), + onDetailsToggle, + loadOlderTurns: vi.fn(), + renderMarkdown: (parent, text) => parent.createDiv({ text }), + renderTextWithWikiLinks: (parent, text) => parent.createDiv({ text }), + copyText: vi.fn(), + canForkItem: () => true, + onForkItem, + })[0]; + + const closedElement = renderMessageBlockElement(closedBlock); + expect(closedElement.querySelector(".codex-panel__copy-message")).not.toBeNull(); + const initialFork = expectPresent(closedElement.querySelector(".codex-panel__fork-message")); + expect(initialFork.getAttribute("aria-label")).toBe("Fork from here"); + initialFork.click(); + expect(onDetailsToggle).toHaveBeenCalledWith("message:fork-actions:a1", true); + + const openBlock = messageStreamBlocks({ + activeThreadId: "thread", + turnLifecycle: idleTurnLifecycle(), + historyCursor: null, + loadingHistory: false, + displayItems: [item], + openDetails: new Set(["message:fork-actions:a1"]), + onDetailsToggle, + loadOlderTurns: vi.fn(), + renderMarkdown: (parent, text) => parent.createDiv({ text }), + renderTextWithWikiLinks: (parent, text) => parent.createDiv({ text }), + copyText: vi.fn(), + canForkItem: () => true, + onForkItem, + })[0]; + + const openElement = renderMessageBlockElement(openBlock); + expect(openElement.querySelector(".codex-panel__copy-message")).toBeNull(); + expect(openElement.querySelector(".codex-panel__fork-and-archive-message")?.getAttribute("aria-label")).toBe( + "Fork and archive", + ); + const openFork = expectPresent(openElement.querySelector(".codex-panel__fork-message")); + expect(openFork.getAttribute("aria-label")).toBe("Fork"); + openFork.click(); + expect(onForkItem).toHaveBeenCalledWith(expect.objectContaining({ id: "a1" }), false); + }); + + it("runs fork and archive from the expanded assistant fork actions", () => { + const onForkItem = vi.fn(); + const item: DisplayItem = { + id: "a1", + kind: "message", + role: "assistant", + text: "answer", + copyText: "answer", + turnId: "turn-1", + markdown: true, + }; + const block = messageStreamBlocks({ + activeThreadId: "thread", + turnLifecycle: idleTurnLifecycle(), + historyCursor: null, + loadingHistory: false, + displayItems: [item], + openDetails: new Set(["message:fork-actions:a1"]), + onDetailsToggle: vi.fn(), + loadOlderTurns: vi.fn(), + renderMarkdown: (parent, text) => parent.createDiv({ text }), + renderTextWithWikiLinks: (parent, text) => parent.createDiv({ text }), + copyText: vi.fn(), + canForkItem: () => true, + onForkItem, + })[0]; + + const element = renderMessageBlockElement(block); + expectPresent(element.querySelector(".codex-panel__fork-and-archive-message")).click(); + + expect(onForkItem).toHaveBeenCalledWith(expect.objectContaining({ id: "a1" }), true); + }); + it("keeps message React actions mounted in the message stream host", () => { const parent = document.createElement("div"); const copyText = vi.fn();