Skip to content

Commit f9581d2

Browse files
vcfgvvcfgv
authored
fix(chat): improve handling of streaming messages and add pending state management (#46)
1 parent ab9b8b6 commit f9581d2

File tree

2 files changed

+65
-11
lines changed

2 files changed

+65
-11
lines changed

‎src/pages/Chat/index.tsx‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,11 @@ export function Chat() {
7272
}
7373

7474
// Extract streaming text for display
75-
const streamText = streamingMessage ? extractText(streamingMessage) : '';
75+
const streamMsg = streamingMessage && typeof streamingMessage === 'object'
76+
? streamingMessage as unknown as { role?: string; content?: unknown; timestamp?: number }
77+
: null;
78+
const streamText = streamMsg ? extractText(streamMsg) : (typeof streamingMessage === 'string' ? streamingMessage : '');
79+
const hasStreamText = streamText.trim().length > 0;
7680

7781
return (
7882
<div className="flex flex-col -m-6" style={{ height: 'calc(100vh - 2.5rem)' }}>
@@ -101,20 +105,20 @@ export function Chat() {
101105
))}
102106

103107
{/* Streaming message */}
104-
{sending && streamText && (
108+
{sending && hasStreamText && (
105109
<ChatMessage
106110
message={{
107111
role: 'assistant',
108-
content: streamingMessage as unknown as string,
109-
timestamp: streamingTimestamp,
112+
content: streamMsg?.content ?? streamText,
113+
timestamp: streamMsg?.timestamp ?? streamingTimestamp,
110114
}}
111115
showThinking={showThinking}
112116
isStreaming
113117
/>
114118
)}
115119

116120
{/* Typing indicator when sending but no stream yet */}
117-
{sending && !streamText && (
121+
{sending && !hasStreamText && (
118122
<TypingIndicator />
119123
)}
120124
</>

‎src/stores/chat.ts‎

Lines changed: 56 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ interface ChatState {
4848
activeRunId: string | null;
4949
streamingText: string;
5050
streamingMessage: unknown | null;
51+
pendingFinal: boolean;
52+
lastUserMessageAt: number | null;
5153

5254
// Sessions
5355
sessions: ChatSession[];
@@ -98,6 +100,25 @@ function isToolOnlyMessage(message: RawMessage | undefined): boolean {
98100
return hasTool && !hasText && !hasNonToolContent;
99101
}
100102

103+
function hasNonToolAssistantContent(message: RawMessage | undefined): boolean {
104+
if (!message) return false;
105+
if (typeof message.content === 'string' && message.content.trim()) return true;
106+
107+
const content = message.content;
108+
if (Array.isArray(content)) {
109+
for (const block of content as ContentBlock[]) {
110+
if (block.type === 'text' && block.text && block.text.trim()) return true;
111+
if (block.type === 'thinking' && block.thinking && block.thinking.trim()) return true;
112+
if (block.type === 'image') return true;
113+
}
114+
}
115+
116+
const msg = message as unknown as Record<string, unknown>;
117+
if (typeof msg.text === 'string' && msg.text.trim()) return true;
118+
119+
return false;
120+
}
121+
101122
// ── Store ────────────────────────────────────────────────────────
102123

103124
export const useChatStore = create<ChatState>((set, get) => ({
@@ -109,6 +130,8 @@ export const useChatStore = create<ChatState>((set, get) => ({
109130
activeRunId: null,
110131
streamingText: '',
111132
streamingMessage: null,
133+
pendingFinal: false,
134+
lastUserMessageAt: null,
112135

113136
sessions: [],
114137
currentSessionKey: 'main',
@@ -182,6 +205,8 @@ export const useChatStore = create<ChatState>((set, get) => ({
182205
streamingMessage: null,
183206
activeRunId: null,
184207
error: null,
208+
pendingFinal: false,
209+
lastUserMessageAt: null,
185210
});
186211
// Load history for new session
187212
get().loadHistory();
@@ -199,6 +224,8 @@ export const useChatStore = create<ChatState>((set, get) => ({
199224
streamingMessage: null,
200225
activeRunId: null,
201226
error: null,
227+
pendingFinal: false,
228+
lastUserMessageAt: null,
202229
});
203230
// Reload sessions list to include the new one after first message
204231
get().loadSessions();
@@ -222,6 +249,18 @@ export const useChatStore = create<ChatState>((set, get) => ({
222249
const rawMessages = Array.isArray(data.messages) ? data.messages as RawMessage[] : [];
223250
const thinkingLevel = data.thinkingLevel ? String(data.thinkingLevel) : null;
224251
set({ messages: rawMessages, thinkingLevel, loading: false });
252+
const { pendingFinal, lastUserMessageAt } = get();
253+
if (pendingFinal) {
254+
const recentAssistant = [...rawMessages].reverse().find((msg) => {
255+
if (msg.role !== 'assistant') return false;
256+
if (!hasNonToolAssistantContent(msg)) return false;
257+
if (lastUserMessageAt && msg.timestamp && msg.timestamp < lastUserMessageAt) return false;
258+
return true;
259+
});
260+
if (recentAssistant) {
261+
set({ sending: false, activeRunId: null, pendingFinal: false });
262+
}
263+
}
225264
} else {
226265
set({ messages: [], loading: false });
227266
}
@@ -252,6 +291,8 @@ export const useChatStore = create<ChatState>((set, get) => ({
252291
error: null,
253292
streamingText: '',
254293
streamingMessage: null,
294+
pendingFinal: false,
295+
lastUserMessageAt: userMsg.timestamp ?? null,
255296
}));
256297

257298
try {
@@ -295,7 +336,7 @@ export const useChatStore = create<ChatState>((set, get) => ({
295336

296337
abortRun: async () => {
297338
const { currentSessionKey } = get();
298-
set({ sending: false, streamingText: '', streamingMessage: null });
339+
set({ sending: false, streamingText: '', streamingMessage: null, pendingFinal: false, lastUserMessageAt: null });
299340

300341
try {
301342
await window.electron.ipcRenderer.invoke(
@@ -331,6 +372,7 @@ export const useChatStore = create<ChatState>((set, get) => ({
331372
const finalMsg = event.message as RawMessage | undefined;
332373
if (finalMsg) {
333374
const toolOnly = isToolOnlyMessage(finalMsg);
375+
const hasOutput = hasNonToolAssistantContent(finalMsg);
334376
const msgId = finalMsg.id || (toolOnly ? `run-${runId}-tool-${Date.now()}` : `run-${runId}`);
335377
set((s) => {
336378
// Check if message already exists (prevent duplicates)
@@ -340,11 +382,13 @@ export const useChatStore = create<ChatState>((set, get) => ({
340382
return toolOnly ? {
341383
streamingText: '',
342384
streamingMessage: null,
385+
pendingFinal: true,
343386
} : {
344387
streamingText: '',
345388
streamingMessage: null,
346-
sending: false,
347-
activeRunId: null,
389+
sending: hasOutput ? false : s.sending,
390+
activeRunId: hasOutput ? null : s.activeRunId,
391+
pendingFinal: hasOutput ? false : true,
348392
};
349393
}
350394
return toolOnly ? {
@@ -355,6 +399,7 @@ export const useChatStore = create<ChatState>((set, get) => ({
355399
}],
356400
streamingText: '',
357401
streamingMessage: null,
402+
pendingFinal: true,
358403
} : {
359404
messages: [...s.messages, {
360405
...finalMsg,
@@ -363,13 +408,14 @@ export const useChatStore = create<ChatState>((set, get) => ({
363408
}],
364409
streamingText: '',
365410
streamingMessage: null,
366-
sending: false,
367-
activeRunId: null,
411+
sending: hasOutput ? false : s.sending,
412+
activeRunId: hasOutput ? null : s.activeRunId,
413+
pendingFinal: hasOutput ? false : true,
368414
};
369415
});
370416
} else {
371417
// No message in final event - reload history to get complete data
372-
set({ streamingText: '', streamingMessage: null, sending: false, activeRunId: null });
418+
set({ streamingText: '', streamingMessage: null, pendingFinal: true });
373419
get().loadHistory();
374420
}
375421
break;
@@ -382,6 +428,8 @@ export const useChatStore = create<ChatState>((set, get) => ({
382428
activeRunId: null,
383429
streamingText: '',
384430
streamingMessage: null,
431+
pendingFinal: false,
432+
lastUserMessageAt: null,
385433
});
386434
break;
387435
}
@@ -391,6 +439,8 @@ export const useChatStore = create<ChatState>((set, get) => ({
391439
activeRunId: null,
392440
streamingText: '',
393441
streamingMessage: null,
442+
pendingFinal: false,
443+
lastUserMessageAt: null,
394444
});
395445
break;
396446
}

0 commit comments

Comments
 (0)