Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 40 additions & 12 deletions chat2db-community-client/src/blocks/AI/index.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ import { buildUserMessageNavigationItems } from './messageNavigation';
import { Pencil } from 'lucide-react';
import MessageNavigationRail from './components/MessageNavigationRail';
import InlineRenameInput from '@/components/InlineRenameInput';
import { AiSessionRequestCoordinator, type AiSessionRequestOwner } from './sessionRequestCoordinator';

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

const handleNewChat = useCallback(() => {
const newSessionOwner = sessionRequestCoordinatorRef.current.beginNewSession();
stop();
setSessionLoading(false);
setAutoFollow(true);
chatInputRef.current?.resetAttachments();
pendingViewportAnchorRef.current = null;
Expand Down Expand Up @@ -1280,6 +1284,7 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
clearChatIdFromPath();
}
onSessionChange?.();
return newSessionOwner;
}, [isPanel, clearChatIdFromPath, onSessionChange, stop]);

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

const handleLoadSessionById = useCallback(
async (sessionId: string, title?: string) => {
const loadOwner = sessionRequestCoordinatorRef.current.beginSessionLoad(sessionId);
const isGenerating = statusRef.current === SSERequestStatus.LOADING;
if (isGenerating) {
const activeSessionId = currentSessionIdRef.current || '';
Expand Down Expand Up @@ -1417,12 +1423,18 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
currentSessionTitleRef.current = inProgressSession.title;
}
inProgressSessionRef.current = null;
if (sessionRequestCoordinatorRef.current.finishSessionLoad(loadOwner)) {
setSessionLoading(false);
}
return;
}

setSessionLoading(true);
try {
const msgs = (await aiStreamService.getChatMessages({ sessionId })) || [];
if (!sessionRequestCoordinatorRef.current.isCurrent(loadOwner)) {
return;
}
const chatItems: IChatItem[] = msgs.map((m: IChatMessage) => ({
id: m.id,
role: m.role as ChatRole,
Expand Down Expand Up @@ -1465,6 +1477,9 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
if (!title) {
try {
const sessions = (await aiStreamService.getChatSessions(undefined as void)) || [];
if (!sessionRequestCoordinatorRef.current.isCurrent(loadOwner)) {
return;
}
const found = sessions.find((s) => s.id === sessionId);
setCurrentSessionTitle(found?.title || '');
currentSessionTitleRef.current = found?.title || '';
Expand All @@ -1473,9 +1488,13 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
}
}
} catch {
feedback.error(i18n('stream.error.loadSessionMessages'));
if (sessionRequestCoordinatorRef.current.isCurrent(loadOwner)) {
feedback.error(i18n('stream.error.loadSessionMessages'));
}
} finally {
setSessionLoading(false);
if (sessionRequestCoordinatorRef.current.finishSessionLoad(loadOwner)) {
setSessionLoading(false);
}
}
},
[onSessionChange, stop],
Expand Down Expand Up @@ -1543,7 +1562,7 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
// Send a message.

const handleSend = useCallback(
async (params: SendParams) => {
async (params: SendParams, sessionOwner?: AiSessionRequestOwner) => {
const content = (params.input || '').trim();
if (!content) return;

Expand All @@ -1557,6 +1576,14 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
feedback.warning(i18n('stream.warning.invalidModel'));
return;
}
const sessionContext = sessionRequestCoordinatorRef.current.resolveSendContext(
sessionOwner,
currentSessionIdRef.current,
messagesRef.current,
);
if (!sessionContext) {
return;
}

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

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

console.log('[AI stream] sending request', {
inputPreview: content.slice(0, 200),
sessionId: currentSessionId || undefined,
sessionId: sessionContext.sessionId,
dataSourceId: params.dataSourceId,
databaseName: params.databaseName,
schemaName: params.schemaName,
Expand All @@ -1631,10 +1661,10 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
})),
});

const isNewSession = !currentSessionId;
const isNewSession = !sessionContext.sessionId;
const requestPromise = request({
input: content,
sessionId: currentSessionId || undefined,
sessionId: sessionContext.sessionId,
history: historyPayload,
enableTools: true,
...modelRequestPayload,
Expand All @@ -1658,9 +1688,7 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
await requestPromise;
},
[
currentSessionId,
isCurrentRoundOverflowingViewport,
messages,
modelOptionMap,
selectedModel?.value,
request,
Expand All @@ -1678,10 +1706,10 @@ export default function AI({ variant = 'page', onTableClick, onPinSql, onSession
const params = (e as CustomEvent).detail as SendParams;
if (params) {
// Start a new conversation before sending to avoid mixing old context.
handleNewChat();
const newSessionOwner = handleNewChat();
// Wait for handleNewChat state cleanup before sending.
setTimeout(() => {
handleSend(params);
handleSend(params, newSessionOwner);
}, 0);
}
};
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import assert from 'node:assert/strict';
import { existsSync, readFileSync, readdirSync, statSync } from 'node:fs';
import path from 'node:path';
import './sessionRequestCoordinator.test';

const sourceRoot = path.resolve('src');
const retiredFiles = [
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
import assert from 'node:assert/strict';
import { AiSessionRequestCoordinator } from './sessionRequestCoordinator';

function deferred<T>() {
let resolve!: (value: T) => void;
const promise = new Promise<T>((next) => {
resolve = next;
});
return { promise, resolve };
}

async function run() {
const externalCoordinator = new AiSessionRequestCoordinator();
const externalNewSessionOwner = externalCoordinator.beginNewSession();
const externalSendWithOldSession = externalCoordinator.resolveSendContext(
externalNewSessionOwner,
'stale-session',
[{ role: 'user', content: 'stale history' }],
);
assert.deepEqual(
externalSendWithOldSession,
{ sessionId: undefined, history: [] },
'an external new-session send must not inherit the render closure session',
);
const externalSendWithOldHistory = externalCoordinator.resolveSendContext(
externalNewSessionOwner,
null,
[{ role: 'user', content: 'stale history' }],
);
assert.deepEqual(
externalSendWithOldHistory,
{ sessionId: undefined, history: [] },
'an external new-session send must not inherit the render closure history',
);

const coordinator = new AiSessionRequestCoordinator();
const committedSessions: string[] = [];
let loading = false;

const loadSession = async (sessionId: string, response: Promise<string>) => {
const owner = coordinator.beginSessionLoad(sessionId);
loading = true;
try {
const resolvedSession = await response;
if (coordinator.isCurrent(owner)) {
committedSessions.push(resolvedSession);
}
} finally {
if (coordinator.finishSessionLoad(owner)) {
loading = false;
}
}
};

const first = deferred<string>();
const second = deferred<string>();
const firstLoad = loadSession('session-a', first.promise);
const secondLoad = loadSession('session-b', second.promise);

first.resolve('session-a');
await firstLoad;
assert.equal(loading, true, 'an obsolete finally must not clear the latest loading state');
assert.deepEqual(committedSessions, [], 'an obsolete response must not commit while the latest request is pending');

second.resolve('session-b');
await secondLoad;
assert.equal(loading, false);
assert.deepEqual(committedSessions, ['session-b']);

const reverseCoordinator = new AiSessionRequestCoordinator();
const reverseCommits: string[] = [];
const resolveReverseLoad = async (sessionId: string, response: Promise<string>) => {
const owner = reverseCoordinator.beginSessionLoad(sessionId);
const resolvedSession = await response;
if (reverseCoordinator.isCurrent(owner)) {
reverseCommits.push(resolvedSession);
}
reverseCoordinator.finishSessionLoad(owner);
};
const reverseFirst = deferred<string>();
const reverseSecond = deferred<string>();
const reverseFirstLoad = resolveReverseLoad('session-a', reverseFirst.promise);
const reverseSecondLoad = resolveReverseLoad('session-b', reverseSecond.promise);
reverseSecond.resolve('session-b');
await reverseSecondLoad;
reverseFirst.resolve('session-a');
await reverseFirstLoad;
assert.deepEqual(reverseCommits, ['session-b'], 'B must remain visible when A resolves after B');

const pendingLoad = coordinator.beginSessionLoad('session-c');
const newSessionOwner = coordinator.beginNewSession();
assert.equal(coordinator.isCurrent(pendingLoad), false, 'new chat must invalidate an outstanding session load');
assert.equal(coordinator.finishSessionLoad(pendingLoad), false);
assert.equal(coordinator.isCurrent(newSessionOwner), true);

const nextLoad = coordinator.beginSessionLoad('session-d');
assert.equal(coordinator.resolveSendContext(newSessionOwner, 'stale-session', ['stale']), null);
assert.equal(coordinator.isCurrent(nextLoad), true);

console.log('AI session request coordinator tests passed.');
}

run().catch((error) => {
console.error(error);
process.exitCode = 1;
});
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
export interface AiSessionRequestOwner {
generation: number;
sessionId: string | null;
}

export interface AiSessionSendContext<T> {
sessionId: string | undefined;
history: T[];
}

export class AiSessionRequestCoordinator {
private generation = 0;

private currentOwner: AiSessionRequestOwner | null = null;

private activeLoadGeneration: number | null = null;

beginSessionLoad(sessionId: string): AiSessionRequestOwner {
const owner = this.advance(sessionId);
this.activeLoadGeneration = owner.generation;
return owner;
}

beginNewSession(): AiSessionRequestOwner {
const owner = this.advance(null);
this.activeLoadGeneration = null;
return owner;
}

isCurrent(owner: AiSessionRequestOwner): boolean {
return (
this.currentOwner?.generation === owner.generation && this.currentOwner.sessionId === owner.sessionId
);
}

finishSessionLoad(owner: AiSessionRequestOwner): boolean {
if (!this.isCurrent(owner) || this.activeLoadGeneration !== owner.generation) {
return false;
}
this.activeLoadGeneration = null;
return true;
}

resolveSendContext<T>(
owner: AiSessionRequestOwner | undefined,
currentSessionId: string | null,
history: readonly T[],
): AiSessionSendContext<T> | null {
if (owner) {
if (!this.isCurrent(owner)) {
return null;
}
if (owner.sessionId === null) {
return { sessionId: undefined, history: [] };
}
}

const sessionId = currentSessionId || undefined;
return {
sessionId,
history: sessionId ? [] : [...history],
};
}

private advance(sessionId: string | null): AiSessionRequestOwner {
this.generation += 1;
const owner = { generation: this.generation, sessionId };
this.currentOwner = owner;
return owner;
}
}
Loading