Skip to content

Commit 862b39f

Browse files
HandSonicHandSonic
committed
fix(ai): keep latest session request owner
1 parent 144a04e commit 862b39f

4 files changed

Lines changed: 218 additions & 12 deletions

File tree

‎chat2db-community-client/src/blocks/AI/index.tsx‎

Lines changed: 40 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ import { buildUserMessageNavigationItems } from './messageNavigation';
5656
import { Pencil } from 'lucide-react';
5757
import MessageNavigationRail from './components/MessageNavigationRail';
5858
import InlineRenameInput from '@/components/InlineRenameInput';
59+
import { AiSessionRequestCoordinator, type AiSessionRequestOwner } from './sessionRequestCoordinator';
5960

6061
/** detects unclosed text in flowing text ```chart block, return chart and whether there are any unfinished diagrams */
6162
function splitIncompleteChartBlock(text: string): { textBeforeChart: string; hasIncompleteChart: boolean } {
@@ -558,6 +559,7 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
558559
const currentSessionIdRef = useRef<string | null>(null);
559560
const currentSessionTitleRef = useRef('');
560561
const messagesRef = useRef<IChatItem[]>([]);
562+
const sessionRequestCoordinatorRef = useRef(new AiSessionRequestCoordinator());
561563
const currentRoundUserMessageIdRef = useRef<string | null>(null);
562564
const statusRef = useRef<SSERequestStatus>(SSERequestStatus.IDLE);
563565
const inProgressSessionRef = useRef<IInProgressSessionSnapshot | null>(null);
@@ -1235,7 +1237,9 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
12351237
// Start a new conversation.
12361238

12371239
const handleNewChat = useCallback(() => {
1240+
const newSessionOwner = sessionRequestCoordinatorRef.current.beginNewSession();
12381241
stop();
1242+
setSessionLoading(false);
12391243
setAutoFollow(true);
12401244
chatInputRef.current?.resetAttachments();
12411245
pendingViewportAnchorRef.current = null;
@@ -1280,6 +1284,7 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
12801284
clearChatIdFromPath();
12811285
}
12821286
onSessionChange?.();
1287+
return newSessionOwner;
12831288
}, [isPanel, clearChatIdFromPath, onSessionChange, stop]);
12841289

12851290
const startPanelHistoryRename = useCallback((session: IChatSession) => {
@@ -1342,6 +1347,7 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
13421347

13431348
const handleLoadSessionById = useCallback(
13441349
async (sessionId: string, title?: string) => {
1350+
const loadOwner = sessionRequestCoordinatorRef.current.beginSessionLoad(sessionId);
13451351
const isGenerating = statusRef.current === SSERequestStatus.LOADING;
13461352
if (isGenerating) {
13471353
const activeSessionId = currentSessionIdRef.current || '';
@@ -1417,12 +1423,18 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
14171423
currentSessionTitleRef.current = inProgressSession.title;
14181424
}
14191425
inProgressSessionRef.current = null;
1426+
if (sessionRequestCoordinatorRef.current.finishSessionLoad(loadOwner)) {
1427+
setSessionLoading(false);
1428+
}
14201429
return;
14211430
}
14221431

14231432
setSessionLoading(true);
14241433
try {
14251434
const msgs = (await aiStreamService.getChatMessages({ sessionId })) || [];
1435+
if (!sessionRequestCoordinatorRef.current.isCurrent(loadOwner)) {
1436+
return;
1437+
}
14261438
const chatItems: IChatItem[] = msgs.map((m: IChatMessage) => ({
14271439
id: m.id,
14281440
role: m.role as ChatRole,
@@ -1465,6 +1477,9 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
14651477
if (!title) {
14661478
try {
14671479
const sessions = (await aiStreamService.getChatSessions(undefined as void)) || [];
1480+
if (!sessionRequestCoordinatorRef.current.isCurrent(loadOwner)) {
1481+
return;
1482+
}
14681483
const found = sessions.find((s) => s.id === sessionId);
14691484
setCurrentSessionTitle(found?.title || '');
14701485
currentSessionTitleRef.current = found?.title || '';
@@ -1473,9 +1488,13 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
14731488
}
14741489
}
14751490
} catch {
1476-
feedback.error(i18n('stream.error.loadSessionMessages'));
1491+
if (sessionRequestCoordinatorRef.current.isCurrent(loadOwner)) {
1492+
feedback.error(i18n('stream.error.loadSessionMessages'));
1493+
}
14771494
} finally {
1478-
setSessionLoading(false);
1495+
if (sessionRequestCoordinatorRef.current.finishSessionLoad(loadOwner)) {
1496+
setSessionLoading(false);
1497+
}
14791498
}
14801499
},
14811500
[onSessionChange, stop],
@@ -1543,7 +1562,7 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
15431562
// Send a message.
15441563

15451564
const handleSend = useCallback(
1546-
async (params: SendParams) => {
1565+
async (params: SendParams, sessionOwner?: AiSessionRequestOwner) => {
15471566
const content = (params.input || '').trim();
15481567
if (!content) return;
15491568

@@ -1557,6 +1576,14 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
15571576
feedback.warning(i18n('stream.warning.invalidModel'));
15581577
return;
15591578
}
1579+
const sessionContext = sessionRequestCoordinatorRef.current.resolveSendContext(
1580+
sessionOwner,
1581+
currentSessionIdRef.current,
1582+
messagesRef.current,
1583+
);
1584+
if (!sessionContext) {
1585+
return;
1586+
}
15601587

15611588
setStreamTraceEntries([]);
15621589
streamTraceEntriesRef.current = [];
@@ -1601,9 +1628,9 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
16011628
});
16021629

16031630
// Let the backend load history for an existing session; otherwise send local history.
1604-
const historyPayload = currentSessionId
1631+
const historyPayload = sessionContext.sessionId
16051632
? []
1606-
: messages
1633+
: sessionContext.history
16071634
.slice(-MAX_HISTORY_ROUNDS * 2)
16081635
.filter((item) => item.content?.trim())
16091636
.map((item) => ({ role: item.role, content: item.content }));
@@ -1613,10 +1640,13 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
16131640
feedback.warning(i18n('stream.warning.invalidModel'));
16141641
return;
16151642
}
1643+
if (sessionOwner && !sessionRequestCoordinatorRef.current.isCurrent(sessionOwner)) {
1644+
return;
1645+
}
16161646

16171647
console.log('[AI stream] sending request', {
16181648
inputPreview: content.slice(0, 200),
1619-
sessionId: currentSessionId || undefined,
1649+
sessionId: sessionContext.sessionId,
16201650
dataSourceId: params.dataSourceId,
16211651
databaseName: params.databaseName,
16221652
schemaName: params.schemaName,
@@ -1631,10 +1661,10 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
16311661
})),
16321662
});
16331663

1634-
const isNewSession = !currentSessionId;
1664+
const isNewSession = !sessionContext.sessionId;
16351665
const requestPromise = request({
16361666
input: content,
1637-
sessionId: currentSessionId || undefined,
1667+
sessionId: sessionContext.sessionId,
16381668
history: historyPayload,
16391669
enableTools: true,
16401670
...modelRequestPayload,
@@ -1658,9 +1688,7 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
16581688
await requestPromise;
16591689
},
16601690
[
1661-
currentSessionId,
16621691
isCurrentRoundOverflowingViewport,
1663-
messages,
16641692
modelOptionMap,
16651693
selectedModel?.value,
16661694
request,
@@ -1678,10 +1706,10 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
16781706
const params = (e as CustomEvent).detail as SendParams;
16791707
if (params) {
16801708
// Start a new conversation before sending to avoid mixing old context.
1681-
handleNewChat();
1709+
const newSessionOwner = handleNewChat();
16821710
// Wait for handleNewChat state cleanup before sending.
16831711
setTimeout(() => {
1684-
handleSend(params);
1712+
handleSend(params, newSessionOwner);
16851713
}, 0);
16861714
}
16871715
};

‎chat2db-community-client/src/blocks/AI/retiredAiSurface.test.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import assert from 'node:assert/strict';
22
import { existsSync, readFileSync, readdirSync, statSync } from 'node:fs';
33
import path from 'node:path';
4+
import './sessionRequestCoordinator.test';
45

56
const sourceRoot = path.resolve('src');
67
const retiredFiles = [

‎chat2db-community-client/src/blocks/AI/sessionRequestCoordinator.test.ts‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
import assert from 'node:assert/strict';
2+
import { AiSessionRequestCoordinator } from './sessionRequestCoordinator';
3+
4+
function deferred<T>() {
5+
let resolve!: (value: T) => void;
6+
const promise = new Promise<T>((next) => {
7+
resolve = next;
8+
});
9+
return { promise, resolve };
10+
}
11+
12+
async function run() {
13+
const externalCoordinator = new AiSessionRequestCoordinator();
14+
const externalNewSessionOwner = externalCoordinator.beginNewSession();
15+
const externalSendWithOldSession = externalCoordinator.resolveSendContext(
16+
externalNewSessionOwner,
17+
'stale-session',
18+
[{ role: 'user', content: 'stale history' }],
19+
);
20+
assert.deepEqual(
21+
externalSendWithOldSession,
22+
{ sessionId: undefined, history: [] },
23+
'an external new-session send must not inherit the render closure session',
24+
);
25+
const externalSendWithOldHistory = externalCoordinator.resolveSendContext(
26+
externalNewSessionOwner,
27+
null,
28+
[{ role: 'user', content: 'stale history' }],
29+
);
30+
assert.deepEqual(
31+
externalSendWithOldHistory,
32+
{ sessionId: undefined, history: [] },
33+
'an external new-session send must not inherit the render closure history',
34+
);
35+
36+
const coordinator = new AiSessionRequestCoordinator();
37+
const committedSessions: string[] = [];
38+
let loading = false;
39+
40+
const loadSession = async (sessionId: string, response: Promise<string>) => {
41+
const owner = coordinator.beginSessionLoad(sessionId);
42+
loading = true;
43+
try {
44+
const resolvedSession = await response;
45+
if (coordinator.isCurrent(owner)) {
46+
committedSessions.push(resolvedSession);
47+
}
48+
} finally {
49+
if (coordinator.finishSessionLoad(owner)) {
50+
loading = false;
51+
}
52+
}
53+
};
54+
55+
const first = deferred<string>();
56+
const second = deferred<string>();
57+
const firstLoad = loadSession('session-a', first.promise);
58+
const secondLoad = loadSession('session-b', second.promise);
59+
60+
first.resolve('session-a');
61+
await firstLoad;
62+
assert.equal(loading, true, 'an obsolete finally must not clear the latest loading state');
63+
assert.deepEqual(committedSessions, [], 'an obsolete response must not commit while the latest request is pending');
64+
65+
second.resolve('session-b');
66+
await secondLoad;
67+
assert.equal(loading, false);
68+
assert.deepEqual(committedSessions, ['session-b']);
69+
70+
const reverseCoordinator = new AiSessionRequestCoordinator();
71+
const reverseCommits: string[] = [];
72+
const resolveReverseLoad = async (sessionId: string, response: Promise<string>) => {
73+
const owner = reverseCoordinator.beginSessionLoad(sessionId);
74+
const resolvedSession = await response;
75+
if (reverseCoordinator.isCurrent(owner)) {
76+
reverseCommits.push(resolvedSession);
77+
}
78+
reverseCoordinator.finishSessionLoad(owner);
79+
};
80+
const reverseFirst = deferred<string>();
81+
const reverseSecond = deferred<string>();
82+
const reverseFirstLoad = resolveReverseLoad('session-a', reverseFirst.promise);
83+
const reverseSecondLoad = resolveReverseLoad('session-b', reverseSecond.promise);
84+
reverseSecond.resolve('session-b');
85+
await reverseSecondLoad;
86+
reverseFirst.resolve('session-a');
87+
await reverseFirstLoad;
88+
assert.deepEqual(reverseCommits, ['session-b'], 'B must remain visible when A resolves after B');
89+
90+
const pendingLoad = coordinator.beginSessionLoad('session-c');
91+
const newSessionOwner = coordinator.beginNewSession();
92+
assert.equal(coordinator.isCurrent(pendingLoad), false, 'new chat must invalidate an outstanding session load');
93+
assert.equal(coordinator.finishSessionLoad(pendingLoad), false);
94+
assert.equal(coordinator.isCurrent(newSessionOwner), true);
95+
96+
const nextLoad = coordinator.beginSessionLoad('session-d');
97+
assert.equal(coordinator.resolveSendContext(newSessionOwner, 'stale-session', ['stale']), null);
98+
assert.equal(coordinator.isCurrent(nextLoad), true);
99+
100+
console.log('AI session request coordinator tests passed.');
101+
}
102+
103+
run().catch((error) => {
104+
console.error(error);
105+
process.exitCode = 1;
106+
});

‎chat2db-community-client/src/blocks/AI/sessionRequestCoordinator.ts‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
export interface AiSessionRequestOwner {
2+
generation: number;
3+
sessionId: string | null;
4+
}
5+
6+
export interface AiSessionSendContext<T> {
7+
sessionId: string | undefined;
8+
history: T[];
9+
}
10+
11+
export class AiSessionRequestCoordinator {
12+
private generation = 0;
13+
14+
private currentOwner: AiSessionRequestOwner | null = null;
15+
16+
private activeLoadGeneration: number | null = null;
17+
18+
beginSessionLoad(sessionId: string): AiSessionRequestOwner {
19+
const owner = this.advance(sessionId);
20+
this.activeLoadGeneration = owner.generation;
21+
return owner;
22+
}
23+
24+
beginNewSession(): AiSessionRequestOwner {
25+
const owner = this.advance(null);
26+
this.activeLoadGeneration = null;
27+
return owner;
28+
}
29+
30+
isCurrent(owner: AiSessionRequestOwner): boolean {
31+
return (
32+
this.currentOwner?.generation === owner.generation && this.currentOwner.sessionId === owner.sessionId
33+
);
34+
}
35+
36+
finishSessionLoad(owner: AiSessionRequestOwner): boolean {
37+
if (!this.isCurrent(owner) || this.activeLoadGeneration !== owner.generation) {
38+
return false;
39+
}
40+
this.activeLoadGeneration = null;
41+
return true;
42+
}
43+
44+
resolveSendContext<T>(
45+
owner: AiSessionRequestOwner | undefined,
46+
currentSessionId: string | null,
47+
history: readonly T[],
48+
): AiSessionSendContext<T> | null {
49+
if (owner) {
50+
if (!this.isCurrent(owner)) {
51+
return null;
52+
}
53+
if (owner.sessionId === null) {
54+
return { sessionId: undefined, history: [] };
55+
}
56+
}
57+
58+
const sessionId = currentSessionId || undefined;
59+
return {
60+
sessionId,
61+
history: sessionId ? [] : [...history],
62+
};
63+
}
64+
65+
private advance(sessionId: string | null): AiSessionRequestOwner {
66+
this.generation += 1;
67+
const owner = { generation: this.generation, sessionId };
68+
this.currentOwner = owner;
69+
return owner;
70+
}
71+
}

0 commit comments

Comments
 (0)