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,