Merge pull request #162 from RadonX/fix/resume-session-issues

CrazyBoyM · web-flow · commit bf20ec789dc9 · 2026-01-15T20:57:56.000+08:00
fix: restore toolUseResult and user messages on session resume
diff --git 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
@@ -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
@@ -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({
           <MessageSelector
             erroredToolUseIDs={erroredToolUseIDs}
             unresolvedToolUseIDs={unresolvedToolUseIDs}
-            messages={normalizeMessagesForAPI(messages)}
+            messages={messages.filter(
+              (m): m is UserMessage | AssistantMessage =>
+                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
@@ -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
@@ -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
@@ -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
@@ -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,

Original file line number	Diff line number	Diff line change
`@@ -116,6 +116,9 @@ function normalizeLoadedUser(entry: JsonlUserEntry): Message \| null {`
`116`	`116`	`type: 'user',`
`117`	`117`	`uuid: entry.uuid as any,`
`118`	`118`	`message: entry.message as any,`
	`119`	`+ ...(entry.toolUseResult !== undefined`
	`120`	`+ ? { toolUseResult: { data: entry.toolUseResult, resultForAssistant: '' } }`
	`121`	`+ : {}),`
`119`	`122`	`}`
`120`	`123`	`}`
`121`	`124`