From b8e4f2506df8db22e7d49ad65ce09ab0123bd1af Mon Sep 17 00:00:00 2001 From: Radon Co Date: Thu, 15 Jan 2026 00:42:09 -0800 Subject: [PATCH] fix: restore toolUseResult and user messages on session resume This PR fixes three issues with session resume: 1. toolUseResult was not restored when loading sessions from JSONL, causing crashes in FileEditToolUpdatedMessage and OutputLine components 2. User text messages were never persisted to JSONL, causing the 2xESC undo menu to be empty after resume 3. Messages with mixed content (tool_result + text) were incorrectly filtered from the undo menu Changes: - Restore toolUseResult field in normalizeLoadedUser() - Persist user text messages in query() before the query loop - Extract filterUserTextMessagesForUndo() as a pure function - Fix filtering to check if ALL blocks are tool_result, not just first Fixes #161 Co-Authored-By: Claude Sonnet 4.5 --- src/app/query.ts | 15 ++ src/ui/components/MessageSelector.tsx | 12 +- src/ui/screens/REPL.tsx | 6 +- src/utils/messages/core.ts | 16 ++ src/utils/protocol/kodeAgentSessionLoad.ts | 3 + .../messages-normalization-reorder.test.ts | 26 +++ tests/unit/session-load.test.ts | 216 ++++++++++++++++++ 7 files changed, 283 insertions(+), 11 deletions(-) diff --git a/src/app/query.ts b/src/app/query.ts index efef3abe..f85df1e7 100644 --- a/src/app/query.ts +++ b/src/app/query.ts @@ -489,6 +489,21 @@ export async function* query( toolUseContext.options?.persistSession !== false && process.env.NODE_ENV !== 'test' + // Persist the last user message that triggered this query (if it's a text message, not a tool result) + // This ensures user prompts are saved to the session file for resume/undo functionality + if (shouldPersistSession && messages.length > 0) { + const lastMessage = messages[messages.length - 1] + if ( + lastMessage?.type === 'user' && + (typeof lastMessage.message.content === 'string' || + (Array.isArray(lastMessage.message.content) && + lastMessage.message.content.length > 0 && + lastMessage.message.content[0]?.type !== 'tool_result')) + ) { + appendSessionJsonlFromMessage({ message: lastMessage, toolUseContext }) + } + } + for await (const message of queryCore( messages, systemPrompt, diff --git a/src/ui/components/MessageSelector.tsx b/src/ui/components/MessageSelector.tsx index b6dc5660..b735805a 100644 --- a/src/ui/components/MessageSelector.tsx +++ b/src/ui/components/MessageSelector.tsx @@ -8,6 +8,7 @@ import { randomUUID } from 'crypto' import { type Tool } from '@tool' import { createUserMessage, + filterUserTextMessagesForUndo, isEmptyMessageText, isNotEmptyMessage, normalizeMessages, @@ -49,16 +50,7 @@ export function MessageSelector({ const allItems = useMemo( () => [ - ...messages - .filter( - _ => - !( - _.type === 'user' && - Array.isArray(_.message.content) && - _.message.content[0]?.type === 'tool_result' - ), - ) - .filter(_ => _.type !== 'assistant'), + ...filterUserTextMessagesForUndo(messages), { ...createUserMessage(''), uuid: currentUUID } as UserMessage, ], [messages, currentUUID], diff --git a/src/ui/screens/REPL.tsx b/src/ui/screens/REPL.tsx index a508dd47..4b201e36 100644 --- a/src/ui/screens/REPL.tsx +++ b/src/ui/screens/REPL.tsx @@ -38,6 +38,7 @@ import { type BinaryFeedbackResult, type Message as MessageType, type ProgressMessage, + type UserMessage, query, } from '@query' import type { WrappedClient } from '@services/mcpClient' @@ -741,7 +742,10 @@ export function REPL({ + m.type === 'user' || m.type === 'assistant', + )} onSelect={async message => { setIsMessageSelectorVisible(false) diff --git a/src/utils/messages/core.ts b/src/utils/messages/core.ts index 9c05b6f3..a13b3633 100644 --- a/src/utils/messages/core.ts +++ b/src/utils/messages/core.ts @@ -608,6 +608,22 @@ export function isEmptyMessageText(text: string): boolean { text.trim() === NO_CONTENT_MESSAGE ) } + +/** + * Filter messages to get user text messages for the undo menu (2xESC). + * Excludes: + * - Assistant messages + * - User messages that only contain tool_result blocks + */ +export function filterUserTextMessagesForUndo( + messages: (UserMessage | AssistantMessage)[], +): UserMessage[] { + return messages.filter((msg): msg is UserMessage => { + if (msg.type !== 'user') return false + if (!Array.isArray(msg.message.content)) return true + return !msg.message.content.every(block => block.type === 'tool_result') + }) +} const STRIPPED_TAGS = [ 'commit_analysis', 'context', diff --git a/src/utils/protocol/kodeAgentSessionLoad.ts b/src/utils/protocol/kodeAgentSessionLoad.ts index dea8af1c..a2b38153 100644 --- a/src/utils/protocol/kodeAgentSessionLoad.ts +++ b/src/utils/protocol/kodeAgentSessionLoad.ts @@ -116,6 +116,9 @@ function normalizeLoadedUser(entry: JsonlUserEntry): Message | null { type: 'user', uuid: entry.uuid as any, message: entry.message as any, + ...(entry.toolUseResult !== undefined + ? { toolUseResult: { data: entry.toolUseResult, resultForAssistant: '' } } + : {}), } } diff --git a/tests/unit/messages-normalization-reorder.test.ts b/tests/unit/messages-normalization-reorder.test.ts index b99750c3..7cf2643a 100644 --- a/tests/unit/messages-normalization-reorder.test.ts +++ b/tests/unit/messages-normalization-reorder.test.ts @@ -4,6 +4,7 @@ import { createAssistantMessage, createProgressMessage, createUserMessage, + filterUserTextMessagesForUndo, getInProgressToolUseIDs, getUnresolvedToolUseIDs, normalizeMessages, @@ -114,4 +115,29 @@ describe('messages normalization + reordering parity', () => { expect(getUnresolvedToolUseIDs(normalized)).toEqual(new Set(['t1', 't2'])) expect(getInProgressToolUseIDs(normalized)).toEqual(new Set(['t1', 't2'])) }) + + test('filterUserTextMessagesForUndo excludes tool_result-only messages', () => { + const messages = [ + createUserMessage('hello'), + makeToolResult('t1'), + createAssistantMessage('response'), + ] + const result = filterUserTextMessagesForUndo(messages as any) + expect(result).toHaveLength(1) + expect(result[0]!.message.content).toBe('hello') + }) + + test('filterUserTextMessagesForUndo keeps user text after tool_results', () => { + const messages = [ + createUserMessage('first'), + createAssistantMessage('response'), + makeToolResult('t1'), + makeToolResult('t2'), + createUserMessage('second'), + ] + const result = filterUserTextMessagesForUndo(messages as any) + expect(result).toHaveLength(2) + expect(result[0]!.message.content).toBe('first') + expect(result[1]!.message.content).toBe('second') + }) }) diff --git a/tests/unit/session-load.test.ts b/tests/unit/session-load.test.ts index 82590e64..bd690321 100644 --- a/tests/unit/session-load.test.ts +++ b/tests/unit/session-load.test.ts @@ -163,6 +163,222 @@ describe('session loader (projects/*.jsonl)', () => { ) }) + test('loads toolUseResult data from user messages with tool results', () => { + const sessionId = '66666666-6666-6666-6666-666666666666' + const path = getSessionLogFilePath({ cwd: projectDir, sessionId }) + mkdirSync( + join( + configDir, + 'projects', + sanitizeProjectNameForSessionStore(projectDir), + ), + { + recursive: true, + }, + ) + + // Simulate a session with a Bash tool result that has toolUseResult data + const lines = + [ + JSON.stringify({ + type: 'file-history-snapshot', + messageId: 'm1', + snapshot: { + messageId: 'm1', + trackedFileBackups: {}, + timestamp: new Date().toISOString(), + }, + isSnapshotUpdate: false, + }), + JSON.stringify({ + type: 'user', + sessionId, + uuid: 'u1', + message: { role: 'user', content: 'run ls command' }, + }), + JSON.stringify({ + type: 'assistant', + sessionId, + uuid: 'a1', + message: { + id: 'msg1', + model: 'x', + type: 'message', + role: 'assistant', + content: [ + { type: 'text', text: 'Running ls...' }, + { type: 'tool_use', id: 'toolu_bash1', name: 'Bash', input: { command: 'ls' } }, + ], + stop_reason: 'tool_use', + stop_sequence: null, + usage: { input_tokens: 0, output_tokens: 0 }, + }, + }), + // User message with tool_result AND toolUseResult data (as saved by kodeAgentSessionLog) + JSON.stringify({ + type: 'user', + sessionId, + uuid: 'u2', + message: { + role: 'user', + content: [ + { + type: 'tool_result', + tool_use_id: 'toolu_bash1', + is_error: false, + content: 'file1.ts\nfile2.ts', + }, + ], + }, + toolUseResult: { + stdout: 'file1.ts\nfile2.ts', + stderr: '', + exitCode: 0, + interrupted: false, + }, + }), + ].join('\n') + '\n' + writeFileSync(path, lines, 'utf8') + + const messages = loadKodeAgentSessionMessages({ + cwd: projectDir, + sessionId, + }) + + expect(messages.length).toBe(3) + + // Verify the tool result message has toolUseResult restored + const toolResultMsg = messages[2] as any + expect(toolResultMsg.type).toBe('user') + expect(toolResultMsg.toolUseResult).toBeDefined() + expect(toolResultMsg.toolUseResult.data).toEqual({ + stdout: 'file1.ts\nfile2.ts', + stderr: '', + exitCode: 0, + interrupted: false, + }) + }) + + test('loads FileEdit toolUseResult with filePath for UI rendering', () => { + const sessionId = '77777777-7777-7777-7777-777777777777' + const path = getSessionLogFilePath({ cwd: projectDir, sessionId }) + mkdirSync( + join( + configDir, + 'projects', + sanitizeProjectNameForSessionStore(projectDir), + ), + { + recursive: true, + }, + ) + + // Simulate a session with a FileEdit tool result + const lines = + [ + JSON.stringify({ + type: 'file-history-snapshot', + messageId: 'm1', + snapshot: { + messageId: 'm1', + trackedFileBackups: {}, + timestamp: new Date().toISOString(), + }, + isSnapshotUpdate: false, + }), + JSON.stringify({ + type: 'user', + sessionId, + uuid: 'u1', + message: { + role: 'user', + content: [ + { + type: 'tool_result', + tool_use_id: 'toolu_edit1', + is_error: false, + content: 'File edited successfully', + }, + ], + }, + // This is the data shape that FileEditToolUpdatedMessage expects + toolUseResult: { + filePath: '/path/to/file.ts', + structuredPatch: [ + { + oldStart: 1, + oldLines: 1, + newStart: 1, + newLines: 2, + lines: ['-old line', '+new line', '+another line'], + }, + ], + }, + }), + ].join('\n') + '\n' + writeFileSync(path, lines, 'utf8') + + const messages = loadKodeAgentSessionMessages({ + cwd: projectDir, + sessionId, + }) + + expect(messages.length).toBe(1) + + const toolResultMsg = messages[0] as any + expect(toolResultMsg.toolUseResult).toBeDefined() + expect(toolResultMsg.toolUseResult.data.filePath).toBe('/path/to/file.ts') + expect(toolResultMsg.toolUseResult.data.structuredPatch).toHaveLength(1) + }) + + test('handles user messages without toolUseResult gracefully', () => { + const sessionId = '88888888-8888-8888-8888-888888888888' + const path = getSessionLogFilePath({ cwd: projectDir, sessionId }) + mkdirSync( + join( + configDir, + 'projects', + sanitizeProjectNameForSessionStore(projectDir), + ), + { + recursive: true, + }, + ) + + // User message without toolUseResult (plain text message) + const lines = + [ + JSON.stringify({ + type: 'file-history-snapshot', + messageId: 'm1', + snapshot: { + messageId: 'm1', + trackedFileBackups: {}, + timestamp: new Date().toISOString(), + }, + isSnapshotUpdate: false, + }), + JSON.stringify({ + type: 'user', + sessionId, + uuid: 'u1', + message: { role: 'user', content: 'hello' }, + // No toolUseResult field + }), + ].join('\n') + '\n' + writeFileSync(path, lines, 'utf8') + + const messages = loadKodeAgentSessionMessages({ + cwd: projectDir, + sessionId, + }) + + expect(messages.length).toBe(1) + const msg = messages[0] as any + expect(msg.type).toBe('user') + expect(msg.toolUseResult).toBeUndefined() + }) + test('findMostRecentKodeAgentSessionId picks newest jsonl by mtime', () => { const projectRoot = join( configDir,