Skip to content

Commit e669c0e

Browse files
DragonnZhangDragonnZhang
committed
fix(desktop): merge processing session refreshes
1 parent 06fc7de commit e669c0e

5 files changed

Lines changed: 225 additions & 32 deletions

File tree

‎apps/electron/src/renderer/App.tsx‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,11 @@ import { stripMarkdown } from './utils/text'
3434
import { getSessionTitle } from './utils/session'
3535
import { coerceInputText } from './lib/input-text'
3636
import { getSessionsToRefreshAfterStaleReconnect } from './lib/reconnect-recovery'
37-
import { formatSessionLoadFailure, shouldTreatSessionLoadFailureAsTransportFallback } from './lib/session-load'
37+
import {
38+
formatSessionLoadFailure,
39+
mergeSessionRefreshResult,
40+
shouldTreatSessionLoadFailureAsTransportFallback,
41+
} from './lib/session-load'
3842
import { extractWorkspaceSlugFromPath } from '@craft-agent/shared/utils/workspace-slug'
3943
import { DEFAULT_THINKING_LEVEL } from '@craft-agent/shared/agent/thinking-levels'
4044
import { initRendererPerf } from './lib/perf'
@@ -506,16 +510,16 @@ export default function App() {
506510
if (!fresh) return 'failed'
507511

508512
const prevSession = store.get(sessionAtomFamily(sessionId))
509-
const preservedStaleMessages = !!prevSession && prevSession.messages.length > 0 && (!fresh.messages || fresh.messages.length === 0)
510-
const nextSession = preservedStaleMessages
511-
? { ...fresh, messages: prevSession.messages }
512-
: fresh
513+
const {
514+
session: nextSession,
515+
preservedExistingMessages,
516+
} = mergeSessionRefreshResult(prevSession, fresh)
513517

514518
clearStreamingState(sessionId)
515519
updateSessionDirect(sessionId, () => nextSession)
516520
syncSessionOptionsFromSession(nextSession)
517521
void reconcilePermissionModeState(sessionId)
518-
return preservedStaleMessages ? 'preserved_stale_messages' : 'refreshed'
522+
return preservedExistingMessages ? 'preserved_stale_messages' : 'refreshed'
519523
} catch (err) {
520524
console.error(`[App] Failed to refresh session ${sessionId}:`, err)
521525
return 'failed'

‎apps/electron/src/renderer/atoms/__tests__/sessions.test.ts‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,40 @@ describe('session message loading atoms', () => {
174174
expect(store.get(loadedSessionsAtom).has(sessionId)).toBe(true)
175175
})
176176

177+
it('does not let a shorter processing response replace existing history', async () => {
178+
const store = createStore()
179+
const sessionId = 'session-1'
180+
181+
globalThis.window = {
182+
electronAPI: {
183+
getSessionMessages: async (id: string) => makeSession({
184+
id,
185+
isProcessing: true,
186+
messages: [
187+
{ ...msg('m3', 'assistant'), content: 'fresh:m3' },
188+
msg('m4', 'assistant'),
189+
],
190+
}),
191+
},
192+
} as unknown as typeof window
193+
194+
store.set(sessionAtomFamily(sessionId), makeSession({
195+
id: sessionId,
196+
isProcessing: true,
197+
messages: [msg('m1'), msg('m2', 'assistant'), msg('m3', 'assistant')],
198+
}))
199+
200+
const result = await store.set(ensureSessionMessagesLoadedAtom, sessionId)
201+
202+
expect(result?.messages.map((message) => [message.id, message.content])).toEqual([
203+
['m1', 'content:m1'],
204+
['m2', 'content:m2'],
205+
['m3', 'fresh:m3'],
206+
['m4', 'content:m4'],
207+
])
208+
expect(store.get(loadedSessionsAtom).has(sessionId)).toBe(false)
209+
})
210+
177211
it('throws when the backend cannot provide messages for a non-empty session', async () => {
178212
const store = createStore()
179213
const sessionId = 'session-1'

‎apps/electron/src/renderer/atoms/sessions.ts‎

Lines changed: 11 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import { atom } from 'jotai'
1212
import type { Getter, Setter } from 'jotai/vanilla'
1313
import { atomFamily } from 'jotai-family'
1414
import type { Session, Message } from '../../shared/types'
15-
import { hasSessionContentHint } from '../lib/session-load'
15+
import { hasSessionContentHint, mergeSessionRefreshResult } from '../lib/session-load'
1616

1717
/**
1818
* Session metadata for list display (lightweight, no messages)
@@ -1112,27 +1112,10 @@ async function loadSessionMessages(
11121112
// so they must be explicitly merged here to be available after app restart.
11131113
const existingSession = get(sessionAtomFamily(sessionId))
11141114
const existingTitle = existingSession?.name ?? existingMeta?.name
1115-
const preservedStaleMessages = !!existingSession
1116-
&& existingSession.messages.length > 0
1117-
&& (!loadedSession.messages || loadedSession.messages.length === 0)
1118-
1119-
const mergedSession = existingSession
1115+
const candidateSession = existingSession
11201116
? {
11211117
...existingSession,
1122-
// CRITICAL: Don't clobber messages if session is actively streaming
1123-
// AND already has messages in the atom. Streaming events update the atom
1124-
// directly and may contain messages the IPC response doesn't know about
1125-
// (race window between IPC request and response).
1126-
// The `messages.length > 0` guard ensures Cmd+R reload works: after reload,
1127-
// the atom starts with messages=[] from getSessions(), so IPC response
1128-
// (which has full history from main process memory) must be used.
1129-
// Also guard against sleep/wake edge case: the server may return
1130-
// empty messages if the session subprocess hasn't finished lazy-loading.
1131-
messages: preservedStaleMessages
1132-
? existingSession.messages
1133-
: existingSession.isProcessing && existingSession.messages.length > 0
1134-
? existingSession.messages
1135-
: loadedSession.messages,
1118+
messages: loadedSession.messages,
11361119
availableCommands: loadedSession.availableCommands ?? existingSession.availableCommands,
11371120
availableSkills: loadedSession.availableSkills ?? existingSession.availableSkills,
11381121
availableSkillDetails: loadedSession.availableSkillDetails ?? existingSession.availableSkillDetails,
@@ -1143,6 +1126,10 @@ async function loadSessionMessages(
11431126
: loadedSession.name || !existingTitle
11441127
? loadedSession
11451128
: { ...loadedSession, name: existingTitle }
1129+
const {
1130+
session: mergedSession,
1131+
preservedExistingMessages,
1132+
} = mergeSessionRefreshResult(existingSession, candidateSession)
11461133
set(sessionAtomFamily(sessionId), mergedSession)
11471134

11481135
// Update only lastFinalMessageId in metadata (now computable from loaded messages).
@@ -1161,10 +1148,10 @@ async function loadSessionMessages(
11611148
}
11621149
}
11631150

1164-
// Mark as loaded only when we received a fresh payload.
1165-
// If we had to preserve stale in-memory messages because the backend returned
1166-
// an empty array during lazy-load recovery, keep the session reloadable.
1167-
if (!preservedStaleMessages) {
1151+
// Mark as loaded only when we received a fresh full payload. If we had to
1152+
// preserve existing messages because the backend returned an empty or short
1153+
// processing snapshot, keep the session reloadable.
1154+
if (!preservedExistingMessages) {
11681155
markSessionMessagesLoaded(get, set, sessionId)
11691156
}
11701157

‎apps/electron/src/renderer/lib/__tests__/session-load.test.ts‎

Lines changed: 104 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
import { describe, expect, it } from 'bun:test'
2-
import type { TransportConnectionState } from '../../../shared/types'
2+
import type { Message, Session, TransportConnectionState } from '../../../shared/types'
33
import {
44
formatSessionLoadFailure,
55
hasSessionContentHint,
6+
mergeSessionRefreshResult,
67
shouldShowMissingSessionState,
78
shouldShowForegroundMessageLoading,
89
shouldTreatSessionLoadFailureAsTransportFallback,
@@ -19,6 +20,27 @@ function createState(overrides?: Partial<TransportConnectionState>): TransportCo
1920
}
2021
}
2122

23+
function message(id: string, content = id): Message {
24+
return {
25+
id,
26+
role: 'assistant',
27+
content,
28+
timestamp: Date.now(),
29+
}
30+
}
31+
32+
function session(overrides: Partial<Session>): Session {
33+
return {
34+
id: 'session-1',
35+
workspaceId: 'workspace-1',
36+
workspaceName: 'Workspace',
37+
messages: [],
38+
isProcessing: false,
39+
lastMessageAt: Date.now(),
40+
...overrides,
41+
} as Session
42+
}
43+
2244
describe('shouldTreatSessionLoadFailureAsTransportFallback', () => {
2345
it('returns true for remote reconnecting state', () => {
2446
expect(shouldTreatSessionLoadFailureAsTransportFallback(
@@ -134,3 +156,84 @@ describe('hasSessionContentHint', () => {
134156
expect(hasSessionContentHint({ messageCount: 0 })).toBe(false)
135157
})
136158
})
159+
160+
describe('mergeSessionRefreshResult', () => {
161+
it('keeps existing history when a processing refresh returns an empty message list', () => {
162+
const existing = session({
163+
messages: [message('m1'), message('m2')],
164+
isProcessing: true,
165+
})
166+
const fresh = session({
167+
messages: [],
168+
isProcessing: true,
169+
})
170+
171+
const result = mergeSessionRefreshResult(existing, fresh)
172+
173+
expect(result.preservedExistingMessages).toBe(true)
174+
expect(result.session.messages.map(m => m.id)).toEqual(['m1', 'm2'])
175+
})
176+
177+
it('merges a shorter processing snapshot into existing history', () => {
178+
const existing = session({
179+
messages: [message('m1'), message('m2', 'old'), message('m3')],
180+
isProcessing: true,
181+
})
182+
const fresh = session({
183+
messages: [message('m2', 'fresh'), message('m4')],
184+
isProcessing: true,
185+
})
186+
187+
const result = mergeSessionRefreshResult(existing, fresh)
188+
189+
expect(result.preservedExistingMessages).toBe(true)
190+
expect(result.session.messages.map(m => [m.id, m.content])).toEqual([
191+
['m1', 'm1'],
192+
['m2', 'fresh'],
193+
['m3', 'm3'],
194+
['m4', 'm4'],
195+
])
196+
})
197+
198+
it('does not replace longer streaming text with a shorter processing snapshot', () => {
199+
const existing = session({
200+
messages: [
201+
{
202+
...message('m1', 'hello from renderer'),
203+
isStreaming: true,
204+
},
205+
],
206+
isProcessing: true,
207+
})
208+
const fresh = session({
209+
messages: [
210+
{
211+
...message('m1', 'hello'),
212+
isStreaming: true,
213+
},
214+
],
215+
isProcessing: true,
216+
})
217+
218+
const result = mergeSessionRefreshResult(existing, fresh)
219+
220+
expect(result.preservedExistingMessages).toBe(true)
221+
expect(result.session.messages[0]?.content).toBe('hello from renderer')
222+
})
223+
224+
it('trusts shorter completed refreshes as authoritative', () => {
225+
const existing = session({
226+
messages: [message('m1'), message('m2'), message('m3')],
227+
isProcessing: false,
228+
})
229+
const fresh = session({
230+
messages: [message('m1')],
231+
isProcessing: false,
232+
})
233+
234+
const result = mergeSessionRefreshResult(existing, fresh)
235+
236+
expect(result.preservedExistingMessages).toBe(false)
237+
expect(result.session.messages.map(m => m.id)).toEqual(['m1'])
238+
})
239+
})

‎apps/electron/src/renderer/lib/session-load.ts‎

Lines changed: 66 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import type { TransportConnectionState } from '../../shared/types'
1+
import type { Message, Session, TransportConnectionState } from '../../shared/types'
22

33
export interface SessionContentHint {
44
name?: string
@@ -64,3 +64,68 @@ export function formatSessionLoadFailure(error: unknown): string {
6464
if (typeof error === 'string' && error.trim()) return error
6565
return 'Unknown error'
6666
}
67+
68+
function mergeMessageById(existingMessage: Message, freshMessage: Message): Message {
69+
if (
70+
existingMessage.role === 'assistant'
71+
&& freshMessage.role === 'assistant'
72+
&& existingMessage.isStreaming
73+
&& freshMessage.isStreaming
74+
&& typeof existingMessage.content === 'string'
75+
&& typeof freshMessage.content === 'string'
76+
&& freshMessage.content.length < existingMessage.content.length
77+
) {
78+
return existingMessage
79+
}
80+
81+
return freshMessage
82+
}
83+
84+
function mergeMessagesById(existingMessages: Message[], freshMessages: Message[]): Message[] {
85+
const freshById = new Map(freshMessages.map(message => [message.id, message]))
86+
const seen = new Set<string>()
87+
const merged = existingMessages.map(message => {
88+
seen.add(message.id)
89+
const freshMessage = freshById.get(message.id)
90+
return freshMessage ? mergeMessageById(message, freshMessage) : message
91+
})
92+
93+
for (const message of freshMessages) {
94+
if (!seen.has(message.id)) merged.push(message)
95+
}
96+
97+
return merged
98+
}
99+
100+
export function mergeSessionRefreshResult(
101+
existingSession: Session | null | undefined,
102+
freshSession: Session,
103+
): { session: Session; preservedExistingMessages: boolean } {
104+
if (!existingSession || existingSession.messages.length === 0) {
105+
return { session: freshSession, preservedExistingMessages: false }
106+
}
107+
108+
const freshMessages = freshSession.messages ?? []
109+
if (freshMessages.length === 0) {
110+
return {
111+
session: { ...freshSession, messages: existingSession.messages },
112+
preservedExistingMessages: true,
113+
}
114+
}
115+
116+
const shouldMergeProcessingSnapshot =
117+
freshMessages.length <= existingSession.messages.length
118+
&& (freshSession.isProcessing || existingSession.isProcessing)
119+
120+
if (!shouldMergeProcessingSnapshot) {
121+
return { session: freshSession, preservedExistingMessages: false }
122+
}
123+
124+
return {
125+
session: {
126+
...freshSession,
127+
messages: mergeMessagesById(existingSession.messages, freshMessages),
128+
},
129+
preservedExistingMessages: true,
130+
}
131+
}

0 commit comments

Comments
 (0)