From bc93493ed1ba366d77a5247833d288641edad111 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A8=8B=E5=BA=8F=E5=91=98=E9=98=BF=E6=B1=9F=28Relakkes?= =?UTF-8?q?=29?= Date: Mon, 27 Apr 2026 21:48:59 +0800 Subject: [PATCH] Prevent desktop rewind from trimming the wrong turn Desktop rewind previously used the visible user-message index as the primary selector. That can drift from the persisted active chain when hidden or non-rendered user messages exist, so the API now prefers a stable user message id and checks the selected prompt text before mutating transcript or files. Constraint: Desktop UI can hide transcript entries that still exist in the persisted session chain Rejected: Keep index-only rewind | can target an earlier user message after UI/server chain drift Confidence: high Scope-risk: moderate Directive: Do not remove targetUserMessageId or expectedContent guards without reproducing shifted visible-index sessions Tested: bun test src/server/__tests__/sessions.test.ts --timeout 20000 Tested: cd desktop && bun run test MessageList.test.tsx chatStore.test.ts -- --run Tested: desktop/scripts/e2e-rewind-agent-browser.sh Tested: desktop/scripts/e2e-rewind-complex-agent-browser.sh Not-tested: Full desktop Vitest suite still has unrelated locale-sensitive failures --- desktop/scripts/e2e-rewind-agent-browser.sh | 2 + .../e2e-rewind-complex-agent-browser.sh | 250 ++++++++++++++++++ desktop/src/api/sessions.ts | 8 +- .../src/components/chat/MessageList.test.tsx | 73 +++++ desktop/src/components/chat/MessageList.tsx | 6 + desktop/src/stores/chatStore.test.ts | 25 ++ desktop/src/stores/chatStore.ts | 2 +- src/server/__tests__/sessions.test.ts | 94 +++++++ src/server/api/sessions.ts | 16 +- src/server/services/sessionRewindService.ts | 87 +++++- 10 files changed, 544 insertions(+), 19 deletions(-) create mode 100755 desktop/scripts/e2e-rewind-complex-agent-browser.sh diff --git a/desktop/scripts/e2e-rewind-agent-browser.sh b/desktop/scripts/e2e-rewind-agent-browser.sh index d1bb2172..7b2ba1ab 100755 --- a/desktop/scripts/e2e-rewind-agent-browser.sh +++ b/desktop/scripts/e2e-rewind-agent-browser.sh @@ -108,6 +108,8 @@ AB="agent-browser --session ${SESSION_NAME}" ${AB} open "${WEB_URL}" ${AB} wait 2000 +${AB} eval "localStorage.setItem('cc-haha-locale', 'en'); location.reload();" +${AB} wait 2000 ${AB} screenshot "${ARTIFACT_DIR}/01-home.png" >/dev/null ${AB} fill '#sidebar-search' "${UNIQUE_TITLE}" diff --git a/desktop/scripts/e2e-rewind-complex-agent-browser.sh b/desktop/scripts/e2e-rewind-complex-agent-browser.sh new file mode 100755 index 00000000..1739f0f0 --- /dev/null +++ b/desktop/scripts/e2e-rewind-complex-agent-browser.sh @@ -0,0 +1,250 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +DESKTOP_DIR="$ROOT_DIR/desktop" + +if ! command -v agent-browser >/dev/null 2>&1; then + echo "agent-browser is required but was not found in PATH" >&2 + exit 1 +fi + +PORTS="$(bun -e 'import { createServer } from "node:net"; const getPort = () => new Promise((resolve) => { const server = createServer(); server.listen(0, "127.0.0.1", () => { const address = server.address(); const port = typeof address === "object" && address ? address.port : 0; server.close(() => resolve(port)); }); }); const apiPort = await getPort(); const webPort = await getPort(); console.log(`${apiPort} ${webPort}`);')" +API_PORT="${PORTS%% *}" +WEB_PORT="${PORTS##* }" +BASE_URL="http://127.0.0.1:${API_PORT}" +WEB_URL="http://127.0.0.1:${WEB_PORT}/?serverUrl=${BASE_URL}" + +RUN_ID="$(date +%s)-$RANDOM" +SESSION_NAME="cc-haha-rewind-complex-e2e-${RUN_ID}" +ARTIFACT_DIR="$(mktemp -d "/tmp/cc-haha-rewind-complex-e2e-${RUN_ID}-XXXX")" +PROJECT_DIR="${ARTIFACT_DIR}/project" +SERVER_LOG="${ARTIFACT_DIR}/server.log" +WEB_LOG="${ARTIFACT_DIR}/web.log" + +mkdir -p "${PROJECT_DIR}/src" +cat > "${PROJECT_DIR}/src/app.js" <<'EOF' +export const STEP = 'base' + +export function readStep() { + return STEP +} +EOF +cat > "${PROJECT_DIR}/README.md" <<'EOF' +# Complex Rewind E2E Fixture + +Initial README content. +EOF +cat > "${PROJECT_DIR}/package.json" <<'EOF' +{ + "name": "complex-rewind-e2e-fixture", + "private": true, + "type": "module" +} +EOF + +cleanup() { + local exit_code=$? + if [[ -n "${SERVER_PID:-}" ]]; then + kill "${SERVER_PID}" >/dev/null 2>&1 || true + wait "${SERVER_PID}" >/dev/null 2>&1 || true + fi + if [[ -n "${WEB_PID:-}" ]]; then + kill "${WEB_PID}" >/dev/null 2>&1 || true + wait "${WEB_PID}" >/dev/null 2>&1 || true + fi + agent-browser --session "${SESSION_NAME}" close >/dev/null 2>&1 || true + if [[ $exit_code -ne 0 ]]; then + echo "Artifacts kept at: ${ARTIFACT_DIR}" >&2 + else + rm -rf "${ARTIFACT_DIR}" + fi +} +trap cleanup EXIT + +echo "Starting backend on ${BASE_URL}" +( + cd "${ROOT_DIR}" + SERVER_PORT="${API_PORT}" bun run src/server/index.ts --port "${API_PORT}" --host 127.0.0.1 +) >"${SERVER_LOG}" 2>&1 & +SERVER_PID=$! + +echo "Starting web UI on http://127.0.0.1:${WEB_PORT}" +( + cd "${DESKTOP_DIR}" + bun run dev -- --host 127.0.0.1 --port "${WEB_PORT}" +) >"${WEB_LOG}" 2>&1 & +WEB_PID=$! + +wait_for_http() { + local url="$1" + for _ in $(seq 1 120); do + if curl -fsS "$url" >/dev/null 2>&1; then + return 0 + fi + sleep 1 + done + echo "Timed out waiting for ${url}" >&2 + return 1 +} + +try_agent_browser() { + local pid + (agent-browser --session "${SESSION_NAME}" "$@" >/dev/null 2>&1) & + pid=$! + for _ in $(seq 1 5); do + if ! kill -0 "$pid" >/dev/null 2>&1; then + wait "$pid" >/dev/null 2>&1 || true + return 0 + fi + sleep 1 + done + kill "$pid" >/dev/null 2>&1 || true + wait "$pid" >/dev/null 2>&1 || true +} + +try_allow_buttons() { + try_agent_browser find role button click --name "Allow" + try_agent_browser find role button click --name "Allow for session" +} + +wait_for_file_contains() { + local file="$1" + local expected="$2" + local timeout="${3:-180}" + for _ in $(seq 1 "$timeout"); do + if [[ -f "$file" ]] && grep -q "$expected" "$file"; then + return 0 + fi + try_allow_buttons + sleep 2 + done + echo "Timed out waiting for ${file} to contain ${expected}" >&2 + return 1 +} + +wait_for_http "${BASE_URL}/health" +wait_for_http "http://127.0.0.1:${WEB_PORT}" + +SESSION_ID="$(curl -fsS -X POST "${BASE_URL}/api/sessions" \ + -H 'Content-Type: application/json' \ + -d "{\"workDir\":\"${PROJECT_DIR//\"/\\\"}\"}" \ + | bun -e 'const input = await Bun.stdin.text(); const data = JSON.parse(input); console.log(data.sessionId);')" + +UNIQUE_TITLE="Complex Rewind ${RUN_ID}" +curl -fsS -X PATCH "${BASE_URL}/api/sessions/${SESSION_ID}" \ + -H 'Content-Type: application/json' \ + -d "{\"title\":\"${UNIQUE_TITLE//\"/\\\"}\"}" >/dev/null + +APP_FILE="${PROJECT_DIR}/src/app.js" +README_FILE="${PROJECT_DIR}/README.md" +GENERATED_FILE="${PROJECT_DIR}/src/generated.js" +FIRST_PROMPT="Edit src/app.js and replace STEP = 'base' with STEP = 'turn-one'. Do not modify any other file. Reply with DONE when finished." +SECOND_PROMPT="Make exactly these three changes: 1. In src/app.js replace STEP = 'turn-one' with STEP = 'turn-two'. 2. Append a new line to README.md that says Second turn touched README. 3. Create src/generated.js with exactly: export const GENERATED = 'second-turn'. Reply with DONE when finished." + +AB="agent-browser --session ${SESSION_NAME}" + +${AB} open "${WEB_URL}" +${AB} wait 2000 +${AB} eval "localStorage.setItem('cc-haha-locale', 'en'); location.reload();" +${AB} wait 2000 +${AB} screenshot "${ARTIFACT_DIR}/01-home.png" >/dev/null + +${AB} fill '#sidebar-search' "${UNIQUE_TITLE}" +${AB} find text "${UNIQUE_TITLE}" click +${AB} wait 1000 +${AB} screenshot "${ARTIFACT_DIR}/02-session-open.png" >/dev/null + +${AB} fill 'textarea' "${FIRST_PROMPT}" +${AB} find role button click --name "Run" +wait_for_file_contains "${APP_FILE}" "turn-one" 180 +if grep -q "turn-two" "${APP_FILE}" || grep -q "Second turn touched README" "${README_FILE}" || [[ -e "${GENERATED_FILE}" ]]; then + echo "Unexpected second-turn state after first prompt" >&2 + exit 1 +fi +${AB} screenshot "${ARTIFACT_DIR}/03-after-first-turn.png" >/dev/null + +${AB} fill 'textarea' "${SECOND_PROMPT}" +${AB} find role button click --name "Run" +wait_for_file_contains "${APP_FILE}" "turn-two" 180 +wait_for_file_contains "${README_FILE}" "Second turn touched README" 60 +wait_for_file_contains "${GENERATED_FILE}" "second-turn" 60 +${AB} screenshot "${ARTIFACT_DIR}/04-after-second-turn.png" >/dev/null + +${AB} eval "const buttons = [...document.querySelectorAll('button')].filter((node) => node.getAttribute('aria-label') === 'Rewind to here'); if (buttons.length < 2) throw new Error('Expected at least two Rewind buttons, found ' + buttons.length); buttons[buttons.length - 1].click();" +${AB} wait 1500 +${AB} screenshot "${ARTIFACT_DIR}/05-rewind-second-turn-modal.png" >/dev/null +${AB} find role button click --name "Rewind here" + +for _ in $(seq 1 120); do + if grep -q "turn-one" "${APP_FILE}" \ + && ! grep -q "turn-two" "${APP_FILE}" \ + && ! grep -q "Second turn touched README" "${README_FILE}" \ + && [[ ! -e "${GENERATED_FILE}" ]]; then + break + fi + sleep 1 +done + +if ! grep -q "turn-one" "${APP_FILE}" || grep -q "turn-two" "${APP_FILE}"; then + echo "src/app.js was not restored to first-turn state" >&2 + ${AB} screenshot "${ARTIFACT_DIR}/failure-app-restore.png" >/dev/null || true + exit 1 +fi +if grep -q "Second turn touched README" "${README_FILE}"; then + echo "README.md still contains second-turn edit" >&2 + ${AB} screenshot "${ARTIFACT_DIR}/failure-readme-restore.png" >/dev/null || true + exit 1 +fi +if [[ -e "${GENERATED_FILE}" ]]; then + echo "Generated file still exists after rewinding second turn" >&2 + ${AB} screenshot "${ARTIFACT_DIR}/failure-generated-file.png" >/dev/null || true + exit 1 +fi + +${AB} wait 1000 +PREFILL_VALUE="$(${AB} get value 'textarea' | tr -d '\r')" +if [[ "${PREFILL_VALUE}" != "${SECOND_PROMPT}" ]]; then + echo "Composer prefill mismatch after second-turn rewind" >&2 + echo "Expected: ${SECOND_PROMPT}" >&2 + echo "Actual: ${PREFILL_VALUE}" >&2 + ${AB} screenshot "${ARTIFACT_DIR}/failure-prefill.png" >/dev/null || true + exit 1 +fi + +TRANSCRIPT_CHECK="$(curl -fsS "${BASE_URL}/api/sessions/${SESSION_ID}/messages" \ + | FIRST_PROMPT="${FIRST_PROMPT}" SECOND_PROMPT="${SECOND_PROMPT}" bun -e ' +const input = await Bun.stdin.text() +const data = JSON.parse(input) +const userMessages = data.messages.filter((message) => message.type === "user") +const texts = userMessages + .map((message) => { + if (typeof message.content === "string") return message.content + if (Array.isArray(message.content)) { + return message.content + .filter((block) => block && block.type === "text") + .map((block) => block.text || "") + .join("\n") + } + return "" + }) +console.log(JSON.stringify({ + total: data.messages.length, + userCount: userMessages.length, + hasFirstPrompt: texts.includes(process.env.FIRST_PROMPT), + hasSecondPrompt: texts.includes(process.env.SECOND_PROMPT), +})) +')" +TRANSCRIPT_USER_COUNT="$(printf '%s' "${TRANSCRIPT_CHECK}" | bun -e 'const data = JSON.parse(await Bun.stdin.text()); console.log(data.userCount)')" +TRANSCRIPT_HAS_FIRST="$(printf '%s' "${TRANSCRIPT_CHECK}" | bun -e 'const data = JSON.parse(await Bun.stdin.text()); console.log(data.hasFirstPrompt ? "1" : "0")')" +TRANSCRIPT_HAS_SECOND="$(printf '%s' "${TRANSCRIPT_CHECK}" | bun -e 'const data = JSON.parse(await Bun.stdin.text()); console.log(data.hasSecondPrompt ? "1" : "0")')" +if [[ "${TRANSCRIPT_USER_COUNT}" != "1" || "${TRANSCRIPT_HAS_FIRST}" != "1" || "${TRANSCRIPT_HAS_SECOND}" != "0" ]]; then + echo "Expected rewound session to keep only the first user prompt, got ${TRANSCRIPT_CHECK}" >&2 + exit 1 +fi + +${AB} screenshot "${ARTIFACT_DIR}/06-after-second-turn-rewind.png" >/dev/null + +echo "Complex rewind E2E passed" +echo "API port: ${API_PORT}" +echo "Web port: ${WEB_PORT}" diff --git a/desktop/src/api/sessions.ts b/desktop/src/api/sessions.ts index ef39c789..58763fd0 100644 --- a/desktop/src/api/sessions.ts +++ b/desktop/src/api/sessions.ts @@ -6,6 +6,7 @@ type MessagesResponse = { messages: MessageEntry[] } type CreateSessionResponse = { sessionId: string } export type SessionRewindResponse = { target: { + targetUserMessageId: string userMessageIndex: number userMessageCount: number } @@ -72,7 +73,12 @@ export const sessionsApi = { return api.get<{ commands: Array<{ name: string; description: string }> }>(`/api/sessions/${sessionId}/slash-commands`) }, - rewind(sessionId: string, body: { userMessageIndex: number; dryRun?: boolean }) { + rewind(sessionId: string, body: { + targetUserMessageId?: string + userMessageIndex?: number + expectedContent?: string + dryRun?: boolean + }) { return api.post(`/api/sessions/${sessionId}/rewind`, body, { timeout: 60_000, }) diff --git a/desktop/src/components/chat/MessageList.test.tsx b/desktop/src/components/chat/MessageList.test.tsx index 92f0ba1f..270e86aa 100644 --- a/desktop/src/components/chat/MessageList.test.tsx +++ b/desktop/src/components/chat/MessageList.test.tsx @@ -532,6 +532,7 @@ describe('MessageList nested tool calls', () => { it('opens a rewind preview modal for user messages', async () => { vi.spyOn(sessionsApi, 'rewind').mockResolvedValue({ target: { + targetUserMessageId: 'user-1', userMessageIndex: 0, userMessageCount: 1, }, @@ -570,11 +571,83 @@ describe('MessageList nested tool calls', () => { expect(within(dialog).getByText('回到这一步重做')).toBeTruthy() expect(within(dialog).getByText('src/example.ts')).toBeTruthy() expect(sessionsApi.rewind).toHaveBeenCalledWith(ACTIVE_TAB, { + targetUserMessageId: 'user-1', userMessageIndex: 0, + expectedContent: '回到这一步重做', dryRun: true, }) }) + it('confirms rewind with the selected message id and prompt guard', async () => { + vi.spyOn(sessionsApi, 'rewind').mockResolvedValue({ + target: { + targetUserMessageId: 'user-2', + userMessageIndex: 1, + userMessageCount: 2, + }, + conversation: { + messagesRemoved: 2, + }, + code: { + available: false, + filesChanged: [], + insertions: 0, + deletions: 0, + }, + }) + const reloadHistory = vi.fn().mockResolvedValue(undefined) + const queueComposerPrefill = vi.fn() + + useChatStore.setState({ + reloadHistory, + queueComposerPrefill, + sessions: { + [ACTIVE_TAB]: makeSessionState({ + messages: [ + { + id: 'user-1', + type: 'user_text', + content: '第一段', + timestamp: 1, + }, + { + id: 'assistant-1', + type: 'assistant_text', + content: 'ok', + timestamp: 2, + }, + { + id: 'user-2', + type: 'user_text', + content: '第二段', + timestamp: 3, + }, + ], + }), + }, + }) + + render() + + const buttons = screen.getAllByRole('button', { name: 'Rewind to here' }) + fireEvent.click(buttons[1]!) + const dialog = await screen.findByRole('dialog') + fireEvent.click(within(dialog).getByRole('button', { name: /Rewind here/ })) + + await waitFor(() => { + expect(sessionsApi.rewind).toHaveBeenLastCalledWith(ACTIVE_TAB, { + targetUserMessageId: 'user-2', + userMessageIndex: 1, + expectedContent: '第二段', + }) + }) + expect(reloadHistory).toHaveBeenCalledWith(ACTIVE_TAB) + expect(queueComposerPrefill).toHaveBeenCalledWith(ACTIVE_TAB, { + text: '第二段', + attachments: undefined, + }) + }) + it('shows raw startup details under translated CLI startup errors', () => { useChatStore.setState({ sessions: { diff --git a/desktop/src/components/chat/MessageList.tsx b/desktop/src/components/chat/MessageList.tsx index 3c6cb537..c59e39c1 100644 --- a/desktop/src/components/chat/MessageList.tsx +++ b/desktop/src/components/chat/MessageList.tsx @@ -151,6 +151,7 @@ export function MessageList({ sessionId }: MessageListProps = {}) { const lastSessionIdRef = useRef(resolvedSessionId) const t = useTranslation() const [rewindTarget, setRewindTarget] = useState<{ + messageId: string userMessageIndex: number content: string attachments?: Extract['attachments'] @@ -187,7 +188,9 @@ export function MessageList({ sessionId }: MessageListProps = {}) { void sessionsApi .rewind(resolvedSessionId, { + targetUserMessageId: rewindTarget.messageId, userMessageIndex: rewindTarget.userMessageIndex, + expectedContent: rewindTarget.content, dryRun: true, }) .then((preview) => { @@ -243,7 +246,9 @@ export function MessageList({ sessionId }: MessageListProps = {}) { } const result = await sessionsApi.rewind(resolvedSessionId, { + targetUserMessageId: rewindTarget.messageId, userMessageIndex: rewindTarget.userMessageIndex, + expectedContent: rewindTarget.content, }) await reloadHistory(resolvedSessionId) @@ -340,6 +345,7 @@ export function MessageList({ sessionId }: MessageListProps = {}) { !isMemberSession ? (message, userMessageIndex) => { setRewindTarget({ + messageId: message.id, userMessageIndex, content: message.content, attachments: message.attachments, diff --git a/desktop/src/stores/chatStore.test.ts b/desktop/src/stores/chatStore.test.ts index c8e975e4..03fb897b 100644 --- a/desktop/src/stores/chatStore.test.ts +++ b/desktop/src/stores/chatStore.test.ts @@ -206,6 +206,31 @@ describe('chatStore history mapping', () => { ]) }) + it('preserves source user ids when restoring array-content user prompts', () => { + const messages: MessageEntry[] = [ + { + id: 'user-with-attachment', + type: 'user', + timestamp: '2026-04-06T00:00:00.000Z', + content: [ + { type: 'text', text: '请看这个文件' }, + { type: 'file', name: 'report.md' }, + ], + }, + ] + + const mapped = mapHistoryMessagesToUiMessages(messages) + + expect(mapped).toMatchObject([ + { + id: 'user-with-attachment', + type: 'user_text', + content: '请看这个文件', + attachments: [{ type: 'file', name: 'report.md' }], + }, + ]) + }) + it('keeps parent tool linkage for live tool events', () => { // Initialize the session first useChatStore.setState({ diff --git a/desktop/src/stores/chatStore.ts b/desktop/src/stores/chatStore.ts index 663e4061..01701006 100644 --- a/desktop/src/stores/chatStore.ts +++ b/desktop/src/stores/chatStore.ts @@ -1043,7 +1043,7 @@ export function mapHistoryMessagesToUiMessages( else if (block.type === 'tool_result') uiMessages.push({ id: nextId(), type: 'tool_result', toolUseId: block.tool_use_id ?? '', content: block.content, isError: !!block.is_error, timestamp, parentToolUseId: msg.parentToolUseId }) } if (textParts.length > 0 || attachments.length > 0) { - uiMessages.push({ id: nextId(), type: 'user_text', content: textParts.join('\n'), attachments: attachments.length > 0 ? attachments : undefined, timestamp }) + uiMessages.push({ id: msg.id || nextId(), type: 'user_text', content: textParts.join('\n'), attachments: attachments.length > 0 ? attachments : undefined, timestamp }) } } } diff --git a/src/server/__tests__/sessions.test.ts b/src/server/__tests__/sessions.test.ts index b7bf564c..721df398 100644 --- a/src/server/__tests__/sessions.test.ts +++ b/src/server/__tests__/sessions.test.ts @@ -909,6 +909,100 @@ describe('Sessions API', () => { ]) }) + it('POST /api/sessions/:id/rewind should target the selected message id instead of a shifted visible index', async () => { + const sessionId = 'aaaaaaaa-bbbb-cccc-dddd-ffffffffffff' + const firstUserId = crypto.randomUUID() + const firstAssistantId = crypto.randomUUID() + const hiddenUserId = crypto.randomUUID() + const targetUserId = crypto.randomUUID() + const targetAssistantId = crypto.randomUUID() + + await writeSessionFile('-tmp-api-rewind-id-target', sessionId, [ + makeSnapshotEntry(), + makeUserEntry('first prompt', firstUserId), + { + ...makeAssistantEntry('first reply', firstUserId), + uuid: firstAssistantId, + }, + makeUserEntry( + 'internal status that the main chat hides', + hiddenUserId, + ), + makeUserEntry('second visible prompt', targetUserId), + { + ...makeAssistantEntry('second reply', targetUserId), + uuid: targetAssistantId, + }, + ]) + + const executeRes = await fetch(`${baseUrl}/api/sessions/${sessionId}/rewind`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + userMessageIndex: 1, + targetUserMessageId: targetUserId, + expectedContent: 'second visible prompt', + }), + }) + expect(executeRes.status).toBe(200) + + const executeBody = await executeRes.json() as { + target: { targetUserMessageId: string; userMessageIndex: number } + conversation: { messagesRemoved: number; removedMessageIds: string[] } + } + expect(executeBody.target.targetUserMessageId).toBe(targetUserId) + expect(executeBody.target.userMessageIndex).toBe(2) + expect(executeBody.conversation.messagesRemoved).toBe(2) + expect(executeBody.conversation.removedMessageIds).toEqual([ + targetUserId, + targetAssistantId, + ]) + + const remainingMessages = await service.getSessionMessages(sessionId) + expect(remainingMessages.map((message) => message.id)).toEqual([ + firstUserId, + firstAssistantId, + hiddenUserId, + ]) + }) + + it('POST /api/sessions/:id/rewind should reject an index fallback when the selected prompt no longer matches', async () => { + const sessionId = 'aaaaaaaa-bbbb-cccc-dddd-000000000000' + const firstUserId = crypto.randomUUID() + const hiddenUserId = crypto.randomUUID() + const targetUserId = crypto.randomUUID() + + await writeSessionFile('-tmp-api-rewind-index-guard', sessionId, [ + makeSnapshotEntry(), + makeUserEntry('first prompt', firstUserId), + makeUserEntry( + 'internal status that the main chat hides', + hiddenUserId, + ), + makeUserEntry('second visible prompt', targetUserId), + ]) + + const executeRes = await fetch(`${baseUrl}/api/sessions/${sessionId}/rewind`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + userMessageIndex: 1, + expectedContent: 'second visible prompt', + }), + }) + expect(executeRes.status).toBe(400) + + const body = await executeRes.json() as { message: string } + expect(body.message).toContain('does not match the selected prompt') + + const remainingMessages = await service.getSessionMessages(sessionId) + expect(remainingMessages.map((message) => message.id)).toEqual([ + firstUserId, + hiddenUserId, + targetUserId, + ]) + }) + it('POST /api/sessions/:id/rewind should restore a single edited file', async () => { const sessionId = 'bbbbbbbb-bbbb-cccc-dddd-eeeeeeeeeeee' const workDir = path.join(tmpDir, 'single-file-fixture') diff --git a/src/server/api/sessions.ts b/src/server/api/sessions.ts index 091e9fba..c9a908c4 100644 --- a/src/server/api/sessions.ts +++ b/src/server/api/sessions.ts @@ -20,6 +20,7 @@ import { getSkillDirCommands } from '../../skills/loadSkillsDir.js' import { executeSessionRewind, previewSessionRewind, + type RewindTargetSelector, } from '../services/sessionRewindService.js' export async function handleSessionsApi( @@ -267,20 +268,23 @@ async function getGitInfo(sessionId: string): Promise { } async function rewindSession(req: Request, sessionId: string): Promise { - let body: { userMessageIndex?: number; dryRun?: boolean } + let body: RewindTargetSelector & { dryRun?: boolean } try { - body = (await req.json()) as { userMessageIndex?: number; dryRun?: boolean } + body = (await req.json()) as RewindTargetSelector & { dryRun?: boolean } } catch { throw ApiError.badRequest('Invalid JSON body') } - if (!Number.isInteger(body.userMessageIndex)) { - throw ApiError.badRequest('userMessageIndex (integer) is required') + if ( + (typeof body.targetUserMessageId !== 'string' || body.targetUserMessageId.length === 0) && + !Number.isInteger(body.userMessageIndex) + ) { + throw ApiError.badRequest('targetUserMessageId (string) or userMessageIndex (integer) is required') } const result = body.dryRun - ? await previewSessionRewind(sessionId, body.userMessageIndex) - : await executeSessionRewind(sessionId, body.userMessageIndex) + ? await previewSessionRewind(sessionId, body) + : await executeSessionRewind(sessionId, body) return Response.json(result) } diff --git a/src/server/services/sessionRewindService.ts b/src/server/services/sessionRewindService.ts index 53054b5d..29e1f9a0 100644 --- a/src/server/services/sessionRewindService.ts +++ b/src/server/services/sessionRewindService.ts @@ -25,8 +25,15 @@ type RewindCodePreview = { deletions: number } +export type RewindTargetSelector = { + targetUserMessageId?: string + userMessageIndex?: number + expectedContent?: string +} + export type SessionRewindPreview = { target: { + targetUserMessageId: string userMessageIndex: number userMessageCount: number } @@ -55,9 +62,43 @@ function normalizeDiffStats(diffStats: { } } +function normalizePromptText(text: string): string { + return text.replace(/\r\n/g, '\n').trim() +} + +function extractUserPromptText(content: unknown): string { + if (typeof content === 'string') return content + if (!Array.isArray(content)) return '' + + return content + .flatMap((block) => { + if (!block || typeof block !== 'object') return [] + const record = block as Record + return record.type === 'text' && typeof record.text === 'string' + ? [record.text] + : [] + }) + .join('\n') +} + +function assertExpectedPromptMatches( + targetMessage: { content: unknown }, + expectedContent: string | undefined, +): void { + if (expectedContent === undefined) return + + const actual = normalizePromptText(extractUserPromptText(targetMessage.content)) + const expected = normalizePromptText(expectedContent) + if (actual !== expected) { + throw ApiError.badRequest( + 'The resolved rewind target does not match the selected prompt. Refresh the session and try again.', + ) + } +} + async function resolveRewindTarget( sessionId: string, - userMessageIndex: number, + selector: RewindTargetSelector, ): Promise { const activeMessages = await sessionService.getSessionMessages(sessionId) const userMessages = activeMessages.filter((message) => message.type === 'user') @@ -66,20 +107,42 @@ async function resolveRewindTarget( throw ApiError.badRequest('This session has no user messages to rewind.') } + let targetUserMessage = null as (typeof userMessages)[number] | null + let userMessageIndex = -1 + + if (selector.targetUserMessageId) { + const activeMessage = activeMessages.find( + (message) => message.id === selector.targetUserMessageId, + ) + if (activeMessage) { + if (activeMessage.type !== 'user') { + throw ApiError.badRequest('The selected rewind target is not a user message.') + } + targetUserMessage = activeMessage + userMessageIndex = userMessages.findIndex( + (message) => message.id === activeMessage.id, + ) + } + } + + if (!targetUserMessage && Number.isInteger(selector.userMessageIndex)) { + userMessageIndex = selector.userMessageIndex! + if (userMessageIndex >= 0 && userMessageIndex < userMessages.length) { + targetUserMessage = userMessages[userMessageIndex]! + } + } + if ( - !Number.isInteger(userMessageIndex) || + !targetUserMessage || userMessageIndex < 0 || userMessageIndex >= userMessages.length ) { throw ApiError.badRequest( - `Invalid userMessageIndex ${userMessageIndex}. Expected 0-${userMessages.length - 1}.`, + `Invalid rewind target. Expected targetUserMessageId or userMessageIndex 0-${userMessages.length - 1}.`, ) } - const targetUserMessage = userMessages[userMessageIndex] - if (!targetUserMessage) { - throw ApiError.badRequest('The selected user message could not be resolved.') - } + assertExpectedPromptMatches(targetUserMessage, selector.expectedContent) const activeMessageIndex = activeMessages.findIndex( (message) => message.id === targetUserMessage.id, @@ -293,9 +356,9 @@ async function buildCodePreview( export async function previewSessionRewind( sessionId: string, - userMessageIndex: number, + selector: RewindTargetSelector, ): Promise { - const target = await resolveRewindTarget(sessionId, userMessageIndex) + const target = await resolveRewindTarget(sessionId, selector) const workDir = (conversationService.hasSession(sessionId) ? conversationService.getSessionWorkDir(sessionId) @@ -310,6 +373,7 @@ export async function previewSessionRewind( return { target: { + targetUserMessageId: target.targetUserMessageId, userMessageIndex: target.userMessageIndex, userMessageCount: target.userMessageCount, }, @@ -322,9 +386,9 @@ export async function previewSessionRewind( export async function executeSessionRewind( sessionId: string, - userMessageIndex: number, + selector: RewindTargetSelector, ): Promise { - const target = await resolveRewindTarget(sessionId, userMessageIndex) + const target = await resolveRewindTarget(sessionId, selector) const workDir = (conversationService.hasSession(sessionId) ? conversationService.getSessionWorkDir(sessionId) @@ -381,6 +445,7 @@ export async function executeSessionRewind( return { target: { + targetUserMessageId: target.targetUserMessageId, userMessageIndex: target.userMessageIndex, userMessageCount: target.userMessageCount, },