diff --git a/packages/tui/src/routes/session/index.tsx b/packages/tui/src/routes/session/index.tsx index dd9fd34d07cb..d4afda850bf1 100644 --- a/packages/tui/src/routes/session/index.tsx +++ b/packages/tui/src/routes/session/index.tsx @@ -72,6 +72,7 @@ import { usePathFormatter } from "../../context/path-format" import { LocationProvider } from "../../context/location" import { createSessionRows, type PartRef, type SessionRow } from "./rows" import { switchLabel } from "../../util/model" +import { revertedPrompt } from "../../util/revert-prompt" addDefaultParsers(parsers.parsers) @@ -403,12 +404,12 @@ export function Session() { void (async () => { const boundary = session()?.revert?.messageID const list = messages() - let target: string | undefined + let target: SessionMessageUser | undefined for (let i = list.length - 1; i >= 0; i--) { const message = list[i] if (message.type !== "user" || !message.text.trim()) continue if (boundary && message.id >= boundary) continue - target = message.id + target = message break } if (!target) { @@ -416,11 +417,15 @@ export function Session() { dialog.clear() return } - const error = await sdk.api.session.revertStage({ sessionID: route.sessionID, messageID: target }).then( + const error = await sdk.api.session.revertStage({ sessionID: route.sessionID, messageID: target.id }).then( () => undefined, (error) => error, ) if (error) toast.show({ message: errorMessage(error), variant: "error", duration: 5000 }) + if (!error && prompt) { + const restored = revertedPrompt(prompt.current, target) + if (restored) prompt.set(restored) + } dialog.clear() })() }, diff --git a/packages/tui/src/util/revert-prompt.ts b/packages/tui/src/util/revert-prompt.ts new file mode 100644 index 000000000000..1469aa2b6427 --- /dev/null +++ b/packages/tui/src/util/revert-prompt.ts @@ -0,0 +1,41 @@ +import type { SessionMessageUser } from "@opencode-ai/sdk/v2" +import type { PromptInfo } from "../prompt/history" + +export function revertedPrompt(current: PromptInfo, message: SessionMessageUser): PromptInfo | undefined { + const input = current.input.trim() + if (current.parts.length || (input && !"/undo".startsWith(input))) return + + return { + input: message.text, + parts: [ + ...(message.files ?? []).map((file) => ({ + type: "file" as const, + mime: file.mime, + filename: file.name, + url: file.uri, + source: file.source + ? { + type: "file" as const, + path: file.name ?? file.uri, + text: { + start: file.source.start, + end: file.source.end, + value: file.source.text, + }, + } + : undefined, + })), + ...(message.agents ?? []).map((agent) => ({ + type: "agent" as const, + name: agent.name, + source: agent.source + ? { + start: agent.source.start, + end: agent.source.end, + value: agent.source.text, + } + : undefined, + })), + ], + } +} diff --git a/packages/tui/test/util/revert-prompt.test.ts b/packages/tui/test/util/revert-prompt.test.ts new file mode 100644 index 000000000000..7e9085a6195d --- /dev/null +++ b/packages/tui/test/util/revert-prompt.test.ts @@ -0,0 +1,65 @@ +import { describe, expect, test } from "bun:test" +import type { SessionMessageUser } from "@opencode-ai/sdk/v2" +import { revertedPrompt } from "../../src/util/revert-prompt" + +const message: SessionMessageUser = { + id: "message-1", + type: "user", + text: "Fix the tests", + time: { created: 1 }, + files: [ + { + uri: "file:///repo/test.ts", + mime: "text/typescript", + name: "test.ts", + source: { start: 0, end: 8, text: "@test.ts" }, + }, + ], + agents: [{ name: "review", source: { start: 9, end: 16, text: "@review" } }], +} + +describe("reverted prompt", () => { + test("restores the reverted user message into an empty prompt", () => { + expect(revertedPrompt({ input: "", parts: [] }, message)).toEqual({ + input: "Fix the tests", + parts: [ + { + type: "file", + mime: "text/typescript", + filename: "test.ts", + url: "file:///repo/test.ts", + source: { + type: "file", + path: "test.ts", + text: { start: 0, end: 8, value: "@test.ts" }, + }, + }, + { + type: "agent", + name: "review", + source: { start: 9, end: 16, value: "@review" }, + }, + ], + }) + }) + + test.each(["/", "/u", "/un", "/und", "/undo", " /undo "])("replaces the undo autocomplete query %p", (input) => { + expect(revertedPrompt({ input, parts: [] }, message)?.input).toBe("Fix the tests") + }) + + test("preserves an existing text draft", () => { + expect(revertedPrompt({ input: "Keep this", parts: [] }, message)).toBeUndefined() + }) + + test("preserves an existing attachment draft", () => { + expect( + revertedPrompt( + { + input: "", + parts: [{ type: "agent", name: "build" }], + }, + message, + ), + ).toBeUndefined() + }) +})