mirror of
https://github.com/t8y2/dbx.git
synced 2026-10-21 05:20:02 +00:00
fix(ai): cancel active streams before conversation changes
This commit is contained in:
Melody Ruch
@@ -57,6 +57,8 @@ import { useToast } from "@/composables/useToast";
|
||||
import { useNavigationTargets } from "@/composables/useNavigationTargets";
|
||||
import { buildAiContext, resolveAiDatabaseTarget, resolveAiNamespaceSelection, resolveDefaultAiSchema, runAgentStream, isVectorDbType, isValidActionForMode, defaultActionForMode, type AiAction, type AiAssistantMode, type AiSqlFileContext, type CustomPromptContext } from "@/lib/ai/ai";
|
||||
import { isAiConfigModelCandidate } from "@/lib/ai/aiConfigCandidates";
|
||||
import { deleteConversationWithCancellation, stopAiGenerationWithFallback } from "@/lib/ai/aiConversationLifecycle";
|
||||
import { AiGenerationGuard } from "@/lib/ai/aiGenerationGuard";
|
||||
import { isTauriRuntime } from "@/lib/backend/tauriRuntime";
|
||||
import { addConfiguredAiModel, aiModelOptions } from "@/lib/ai/aiConfigList";
|
||||
import { orderAiConfigsForDisplay } from "@/lib/ai/aiConfigOrdering";
|
||||
@@ -274,11 +276,31 @@ let messageScrollViewport: HTMLElement | null = null;
|
||||
let messageTouchStartY: number | null = null;
|
||||
let lastMessageScrollTop = 0;
|
||||
const STREAM_RENDER_INTERVAL_MS = 33;
|
||||
// How long cancelStream() (the Stop button) waits for the backend to actually
|
||||
// acknowledge a cancellation before forcing the same abandon path clear/switch
|
||||
// uses. See cancelStream() for why the backend RPC alone can't be trusted to
|
||||
// unstick a genuinely hung tool call.
|
||||
const STOP_FORCE_ABANDON_MS = 5000;
|
||||
let assistantDeltaFrame: number | null = null;
|
||||
let lastAssistantFlushAt = 0;
|
||||
let pendingAssistantDelta = "";
|
||||
let pendingAssistantReasoning = "";
|
||||
let pendingAssistantIndex = -1;
|
||||
// Index into `messages.value` of the current generation's assistant placeholder,
|
||||
// mirroring `currentSessionId` (set alongside it in send(), cleared in its finally
|
||||
// and in resetPendingRequestState()). Lets cancelStream()'s forced-abandon path
|
||||
// finalize that specific message — the backend session id alone doesn't identify
|
||||
// it, and abandonInFlightRequest() itself is also used by clear/switch/unmount,
|
||||
// where messages.value is being discarded/replaced anyway so it has no reason to
|
||||
// know about individual messages.
|
||||
let currentAssistantMessageIndex = -1;
|
||||
// Identifies which send() invocation is still allowed to write into `messages`/
|
||||
// `isGenerating`/`currentSessionId` and the delta buffers above.
|
||||
// abandonInFlightRequest() (used by clearMessages()/selectConversation()) invalidates
|
||||
// the active generation so a superseded send() can't corrupt state that now belongs
|
||||
// to a different conversation. See lib/ai/aiGenerationGuard.ts for why this exists
|
||||
// instead of relying on isGenerating/currentSessionId alone.
|
||||
const aiGenerationGuard = new AiGenerationGuard();
|
||||
|
||||
function startEditMessage(visibleIndex: number) {
|
||||
if (isGenerating.value) return;
|
||||
@@ -1727,12 +1749,30 @@ async function send() {
|
||||
}
|
||||
// Acquire the send guard before the first async operation so two rapid
|
||||
// submissions cannot both pass the initial isGenerating check and then
|
||||
// resume into concurrent agent runs.
|
||||
// resume into concurrent agent runs. `myGeneration` is this call's identity:
|
||||
// every mutation of shared state below, once execution has been suspended and
|
||||
// resumed at least once, must check `aiGenerationGuard.isCurrent(myGeneration)`
|
||||
// first, since clearMessages()/selectConversation() can invalidate it out from
|
||||
// under an in-flight send().
|
||||
isGenerating.value = true;
|
||||
const myGeneration = aiGenerationGuard.begin();
|
||||
if (!(await promptTemplateStore.ensureLoaded())) {
|
||||
clearPendingWriteGrant();
|
||||
isGenerating.value = false;
|
||||
toast(t("ai.customInstructionsLoadFailed"), 5000);
|
||||
if (aiGenerationGuard.isCurrent(myGeneration)) {
|
||||
isGenerating.value = false;
|
||||
toast(t("ai.customInstructionsLoadFailed"), 5000);
|
||||
}
|
||||
return;
|
||||
}
|
||||
// Superseded (chat cleared/switched, or a newer send() started) while awaiting
|
||||
// the prompt templates above — bail before touching messages/mentions that now
|
||||
// belong to a different conversation. Also clear the pending write-SQL grant:
|
||||
// it hasn't been read/reset yet (that happens below, right before
|
||||
// runAgentStream()), so a bare return here would leave a previously-confirmed
|
||||
// write grant sitting in the module-scope vars, live to be replayed against
|
||||
// whatever unrelated send() the next conversation issues.
|
||||
if (!aiGenerationGuard.isCurrent(myGeneration)) {
|
||||
clearPendingWriteGrant();
|
||||
return;
|
||||
}
|
||||
// Snapshot the selected custom prompts at send time so later async context loading
|
||||
@@ -1821,15 +1861,29 @@ async function send() {
|
||||
confirmedSchema = undefined;
|
||||
messages.value.push({ role: "assistant", content: "", sourceConnectionName: connection.name });
|
||||
const assistantIdx = messages.value.length - 1;
|
||||
currentAssistantMessageIndex = assistantIdx;
|
||||
const sessionId = uuid();
|
||||
currentSessionId.value = sessionId;
|
||||
const agentEvents: AgentEvent[] = [];
|
||||
try {
|
||||
const sqlFiles = await loadReferencedSqlFiles(selectedSqlFiles);
|
||||
// Superseded while awaiting loadReferencedSqlFiles() above — bail before
|
||||
// paying for buildAiContext() too; it can do real backend/schema work that
|
||||
// would be entirely wasted on an already-abandoned request.
|
||||
if (!aiGenerationGuard.isCurrent(myGeneration)) return;
|
||||
const context = await buildAiContext(tab, connection, {
|
||||
mentionedTables,
|
||||
sqlFiles,
|
||||
});
|
||||
// Superseded while awaiting buildAiContext() above — must bail before ever
|
||||
// calling runAgentStream(), not just before writing its results. Without
|
||||
// this recheck, a clear/switch/unmount that fires during context
|
||||
// preparation invalidates the generation but the request still gets sent to
|
||||
// the backend and starts executing tools/SQL; the best-effort cancel RPC
|
||||
// fired by abandonInFlightRequest() is a no-op here since no session has
|
||||
// been registered with the backend yet (registration happens inside
|
||||
// runAgentStream() itself).
|
||||
if (!aiGenerationGuard.isCurrent(myGeneration)) return;
|
||||
const history: AiMessage[] = messagesForAgentHistory(messages.value.slice(0, -2));
|
||||
await runAgentStream(
|
||||
{
|
||||
@@ -1846,6 +1900,10 @@ async function send() {
|
||||
},
|
||||
history,
|
||||
(event: AgentEvent) => {
|
||||
// Superseded by a clear/switch/new-chat (or a newer send()) — the backend
|
||||
// stream may still be running, but this generation no longer owns any
|
||||
// shared state to write into.
|
||||
if (!aiGenerationGuard.isCurrent(myGeneration)) return;
|
||||
agentEvents.push(event);
|
||||
if (event.type === "text_delta" && event.delta) {
|
||||
appendAssistantDelta(assistantIdx, event.delta);
|
||||
@@ -1883,59 +1941,164 @@ async function send() {
|
||||
customPromptContext,
|
||||
);
|
||||
} catch (e: unknown) {
|
||||
const message = e instanceof Error ? e.message : String(e);
|
||||
messages.value[assistantIdx].content = `${t("ai.requestFailed")}\n\n${translateBackendError(t, message)}`;
|
||||
// A superseded generation's error (including one caused by an
|
||||
// abandonInFlightRequest()-triggered cancellation) must not overwrite a
|
||||
// message that now belongs to a different conversation, or one that no
|
||||
// longer exists in `messages.value`.
|
||||
if (aiGenerationGuard.isCurrent(myGeneration)) {
|
||||
const message = e instanceof Error ? e.message : String(e);
|
||||
const msg = messages.value[assistantIdx];
|
||||
if (msg) msg.content = `${t("ai.requestFailed")}\n\n${translateBackendError(t, message)}`;
|
||||
}
|
||||
} finally {
|
||||
if (assistantDeltaFrame !== null) cancelAnimationFrame(assistantDeltaFrame);
|
||||
flushAssistantDeltas();
|
||||
const msg = messages.value[assistantIdx];
|
||||
if (msg) msg.isThinking = false;
|
||||
isGenerating.value = false;
|
||||
// Render agent tool call steps from agent events (fallback when no real-time steps)
|
||||
if (msg && agentEvents.length > 0 && !msg.agentSteps?.length) {
|
||||
const steps: AiAgentStepItem[] = [];
|
||||
agentEvents.forEach((e, index) => {
|
||||
const step = agentEventToStep(e, index);
|
||||
if (step) upsertAgentStep(steps, step);
|
||||
});
|
||||
if (steps.length) msg.agentSteps = steps;
|
||||
}
|
||||
// Fallback: use aiAgentPlan for backward compatibility
|
||||
if (msg && !msg.agentSteps?.length) {
|
||||
const agentPlan = buildAiAgentPlan({
|
||||
mode: requestedMode,
|
||||
action: requestedAction,
|
||||
instruction: modelInstruction,
|
||||
assistantContent: msg?.content || "",
|
||||
connection: connection,
|
||||
database: tab.database,
|
||||
});
|
||||
if (msg && requestedMode === "agent") msg.agentSteps = buildAiAgentStepItems(agentPlan);
|
||||
if (agentPlan.handoffSql) emit("requestAutoExecuteSql", agentPlan.handoffSql);
|
||||
}
|
||||
currentSessionId.value = "";
|
||||
// Apply deferred context compaction after streaming so assistantIdx stays stable.
|
||||
// Visible chat history is kept for the user; future LLM history starts from this hidden summary.
|
||||
if (pendingCompaction.value) {
|
||||
const { summary, compactedMessages } = pendingCompaction.value;
|
||||
pendingCompaction.value = null;
|
||||
const insertAt = Math.min(1 + compactedMessages, messages.value.length - 1);
|
||||
if (summary) {
|
||||
messages.value.splice(insertAt, 0, {
|
||||
role: "user",
|
||||
content: summary,
|
||||
kind: "contextSummary",
|
||||
// Everything below mutates state (messages, isGenerating, currentSessionId,
|
||||
// the delta buffers) that only the current generation is allowed to touch.
|
||||
// A superseded generation's cleanup is a no-op: abandonInFlightRequest()
|
||||
// already reset isGenerating/currentSessionId/delta buffers synchronously
|
||||
// when it invalidated this generation.
|
||||
// This block CONSUMES this generation's per-request transient state
|
||||
// (applies flushed deltas to the message, splices the compaction summary
|
||||
// into history) rather than just discarding it — see
|
||||
// resetPendingRequestState() below for the abandon-path equivalent that
|
||||
// discards it instead. If you add a new piece of per-request transient
|
||||
// state, it must be handled on both paths.
|
||||
if (aiGenerationGuard.isCurrent(myGeneration)) {
|
||||
if (assistantDeltaFrame !== null) cancelAnimationFrame(assistantDeltaFrame);
|
||||
flushAssistantDeltas();
|
||||
const msg = messages.value[assistantIdx];
|
||||
if (msg) msg.isThinking = false;
|
||||
isGenerating.value = false;
|
||||
// Render agent tool call steps from agent events (fallback when no real-time steps)
|
||||
if (msg && agentEvents.length > 0 && !msg.agentSteps?.length) {
|
||||
const steps: AiAgentStepItem[] = [];
|
||||
agentEvents.forEach((e, index) => {
|
||||
const step = agentEventToStep(e, index);
|
||||
if (step) upsertAgentStep(steps, step);
|
||||
});
|
||||
if (steps.length) msg.agentSteps = steps;
|
||||
}
|
||||
// Fallback: use aiAgentPlan for backward compatibility
|
||||
if (msg && !msg.agentSteps?.length) {
|
||||
const agentPlan = buildAiAgentPlan({
|
||||
mode: requestedMode,
|
||||
action: requestedAction,
|
||||
instruction: modelInstruction,
|
||||
assistantContent: msg?.content || "",
|
||||
connection: connection,
|
||||
database: tab.database,
|
||||
});
|
||||
if (msg && requestedMode === "agent") msg.agentSteps = buildAiAgentStepItems(agentPlan);
|
||||
if (agentPlan.handoffSql) emit("requestAutoExecuteSql", agentPlan.handoffSql);
|
||||
}
|
||||
currentSessionId.value = "";
|
||||
currentAssistantMessageIndex = -1;
|
||||
// Apply deferred context compaction after streaming so assistantIdx stays stable.
|
||||
// Visible chat history is kept for the user; future LLM history starts from this hidden summary.
|
||||
if (pendingCompaction.value) {
|
||||
const { summary, compactedMessages } = pendingCompaction.value;
|
||||
pendingCompaction.value = null;
|
||||
const insertAt = Math.min(1 + compactedMessages, messages.value.length - 1);
|
||||
if (summary) {
|
||||
messages.value.splice(insertAt, 0, {
|
||||
role: "user",
|
||||
content: summary,
|
||||
kind: "contextSummary",
|
||||
});
|
||||
}
|
||||
}
|
||||
persistConversation();
|
||||
scrollToBottom();
|
||||
}
|
||||
persistConversation();
|
||||
scrollToBottom();
|
||||
}
|
||||
}
|
||||
|
||||
// Resolves once `isGenerating` goes false, or after `timeoutMs` — whichever
|
||||
// comes first. Used by cancelStream() to bound how long it waits for the
|
||||
// backend to actually acknowledge a cancellation before forcing it.
|
||||
function waitForGenerationToClear(timeoutMs: number): Promise<void> {
|
||||
return new Promise((resolve) => {
|
||||
if (!isGenerating.value) {
|
||||
resolve();
|
||||
return;
|
||||
}
|
||||
const stopWatch = watch(isGenerating, (value) => {
|
||||
if (value) return;
|
||||
stopWatch();
|
||||
clearTimeout(timer);
|
||||
resolve();
|
||||
});
|
||||
const timer = setTimeout(() => {
|
||||
stopWatch();
|
||||
resolve();
|
||||
}, timeoutMs);
|
||||
});
|
||||
}
|
||||
|
||||
async function cancelStream() {
|
||||
if (currentSessionId.value) {
|
||||
await aiCancelStream(currentSessionId.value).catch(() => {});
|
||||
await stopAiGenerationWithFallback({
|
||||
isGenerating: () => isGenerating.value,
|
||||
currentGeneration: () => aiGenerationGuard.peek(),
|
||||
isGenerationCurrent: (generation) => aiGenerationGuard.isCurrent(generation),
|
||||
currentSessionId: () => currentSessionId.value,
|
||||
cancelSession: (sessionId) => aiCancelStream(sessionId).then(() => undefined),
|
||||
waitForGenerationToClear: () => waitForGenerationToClear(STOP_FORCE_ABANDON_MS),
|
||||
flushPending: () => {
|
||||
if (assistantDeltaFrame !== null) cancelAnimationFrame(assistantDeltaFrame);
|
||||
flushAssistantDeltas();
|
||||
},
|
||||
currentAssistantMessageIndex: () => currentAssistantMessageIndex,
|
||||
messageAt: (index) => messages.value[index],
|
||||
cancelledMessage: () => t("ai.requestCancelled"),
|
||||
abandon: (sessionId) => abandonInFlightRequest(sessionId),
|
||||
persistConversation,
|
||||
});
|
||||
}
|
||||
|
||||
// Neutralizes all per-request transient state that must never survive into a
|
||||
// different generation/conversation. abandonInFlightRequest() calls this to
|
||||
// discard it immediately. send()'s finally does NOT call it — that block must
|
||||
// first CONSUME this state (apply flushed deltas to the message, splice the
|
||||
// compaction summary into history) rather than discard it — but if you add a
|
||||
// new piece of per-request transient state, add its reset here so it can't be
|
||||
// missed the way pendingCompaction was (see PR #6332 review).
|
||||
function resetPendingRequestState() {
|
||||
if (assistantDeltaFrame !== null) {
|
||||
cancelAnimationFrame(assistantDeltaFrame);
|
||||
assistantDeltaFrame = null;
|
||||
}
|
||||
pendingAssistantDelta = "";
|
||||
pendingAssistantReasoning = "";
|
||||
pendingAssistantIndex = -1;
|
||||
pendingCompaction.value = null;
|
||||
}
|
||||
|
||||
// `alreadyCancelledSessionId`: the session id a caller (cancelStream()) has
|
||||
// already sent the backend cancel RPC for, if any — pass it so this function
|
||||
// doesn't fire a second, redundant RPC for the same session. Left undefined
|
||||
// by clear/switch/unmount, which never RPC before calling this.
|
||||
function abandonInFlightRequest(alreadyCancelledSessionId?: string) {
|
||||
// Used when the UI is about to move to a different conversation/transcript
|
||||
// (clear chat, switch conversation, new chat) while a request may still be
|
||||
// in flight. Unlike cancelStream() above, this must reset shared state
|
||||
// synchronously and unconditionally:
|
||||
// - the backend cancel RPC depends on a session id having already been
|
||||
// registered (send() only sets currentSessionId partway through), so it
|
||||
// can be a silent no-op if this fires before that point;
|
||||
// - even when the RPC isn't a no-op, waiting for the backend to actually
|
||||
// stop before resetting isGenerating is exactly what stranded the send
|
||||
// box indefinitely in issue #5941.
|
||||
// Invalidating the generation here makes send()'s remaining event callbacks,
|
||||
// catch, and finally no-ops regardless of what the backend does next, so
|
||||
// they can't write into the array this call is about to replace. See
|
||||
// lib/ai/aiGenerationGuard.ts.
|
||||
const sessionId = currentSessionId.value;
|
||||
aiGenerationGuard.invalidate();
|
||||
isGenerating.value = false;
|
||||
currentSessionId.value = "";
|
||||
currentAssistantMessageIndex = -1;
|
||||
resetPendingRequestState();
|
||||
if (sessionId && sessionId !== alreadyCancelledSessionId) {
|
||||
aiCancelStream(sessionId).catch(() => {});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2019,6 +2182,14 @@ async function exportMessageAsMarkdown(msg: ChatMessage) {
|
||||
}
|
||||
|
||||
function clearMessages() {
|
||||
// If a request is still in flight, abandon it before wiping the transcript it
|
||||
// was writing into. abandonInFlightRequest() invalidates the active generation
|
||||
// synchronously, so the in-flight send()'s callbacks/catch/finally become
|
||||
// no-ops even if the backend cancel RPC itself can't reach a registered
|
||||
// session id yet — otherwise isGenerating would never reset (nothing but
|
||||
// send()'s own finally clears it) and the send box would stay stuck disabled
|
||||
// indefinitely.
|
||||
if (isGenerating.value) abandonInFlightRequest();
|
||||
messages.value = [];
|
||||
conversationId.value = "";
|
||||
historyIndex.value = -1;
|
||||
@@ -2058,6 +2229,11 @@ async function setConversationListOpen(open: boolean) {
|
||||
}
|
||||
|
||||
function selectConversation(conv: AiConversation) {
|
||||
// Same guard as clearMessages(): switching away from an in-flight request must
|
||||
// abandon it first — abandonInFlightRequest() invalidates the generation so
|
||||
// the old send() can't write its deltas/result into this (different)
|
||||
// conversation's messages array once it's assigned below.
|
||||
if (isGenerating.value) abandonInFlightRequest();
|
||||
conversationId.value = conv.id;
|
||||
// Drop the previous conversation's rendered Markdown instead of keeping it until the LRU evicts it.
|
||||
messageRenderer.value.clear();
|
||||
@@ -2075,9 +2251,17 @@ function selectConversation(conv: AiConversation) {
|
||||
}
|
||||
|
||||
async function deleteConversation(id: string) {
|
||||
await deleteAiConversation(id).catch(() => {});
|
||||
conversations.value = conversations.value.filter((c) => c.id !== id);
|
||||
if (conversationId.value === id) clearMessages();
|
||||
await deleteConversationWithCancellation({
|
||||
id,
|
||||
currentConversationId: () => conversationId.value,
|
||||
isGenerating: () => isGenerating.value,
|
||||
abandon: () => abandonInFlightRequest(),
|
||||
deletePersisted: () => deleteAiConversation(id).catch(() => {}),
|
||||
afterDelete: () => {
|
||||
conversations.value = conversations.value.filter((c) => c.id !== id);
|
||||
if (conversationId.value === id) clearMessages();
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function startNewChat() {
|
||||
@@ -2164,7 +2348,13 @@ onUnmounted(() => {
|
||||
if (assistantDeltaFrame !== null) cancelAnimationFrame(assistantDeltaFrame);
|
||||
clearTimeout(mentionTimer);
|
||||
clearEffortMenuCloseTimer();
|
||||
cancelStream();
|
||||
// Must invalidate the generation the same way clearMessages()/selectConversation()
|
||||
// do, not just fire the best-effort cancelStream() RPC: if a request is still
|
||||
// mid-await (context preparation, or the backend hasn't registered a session id
|
||||
// yet) when this component unmounts, cancelStream() alone leaves the generation
|
||||
// current, so the request still starts and its event callback/catch/finally keep
|
||||
// writing into refs this now-unmounted instance's closures still hold.
|
||||
if (isGenerating.value) abandonInFlightRequest();
|
||||
detachMessageScrollListener();
|
||||
// 清理拖拽事件监听，防止内存泄漏
|
||||
document.removeEventListener("mousemove", handleResize);
|
||||
|
||||
+242
@@ -0,0 +1,242 @@
|
||||
import { readFileSync } from "node:fs";
|
||||
import { describe, expect, it } from "vitest";
|
||||
|
||||
// Regression coverage for https://github.com/t8y2/dbx/issues/5941 and the follow-up
|
||||
// race identified on https://github.com/t8y2/dbx/pull/6332:
|
||||
//
|
||||
// When an in-flight AI request is stuck (e.g. a hung MCP tool call / slow provider
|
||||
// response), clicking the clear-chat trash icon, "New Chat", deleting the active
|
||||
// conversation, or switching to a different saved conversation must (a) cancel the
|
||||
// in-flight request so `isGenerating` doesn't stay stranded true forever, and (b)
|
||||
// isolate that in-flight request's async event callbacks / catch / finally from
|
||||
// whatever conversation is active by the time they run — even if the backend cancel
|
||||
// RPC itself is a no-op because the request hadn't registered a session id yet.
|
||||
//
|
||||
// (b) is what `AiGenerationGuard` (lib/ai/aiGenerationGuard.ts) exists for. Its own
|
||||
// spec and aiConversationLifecycle.spec.ts exercise the async ordering directly;
|
||||
// this file pins that AiAssistant.vue wires those tested lifecycle helpers into the
|
||||
// component paths that own the shared state.
|
||||
const source = readFileSync(new URL("../AiAssistant.vue", import.meta.url), "utf8");
|
||||
|
||||
function bodyOf(fnSignature: string): string {
|
||||
const start = source.indexOf(fnSignature);
|
||||
expect(start, `expected to find "${fnSignature}" in AiAssistant.vue`).toBeGreaterThanOrEqual(0);
|
||||
const braceStart = source.indexOf("{", start);
|
||||
let depth = 0;
|
||||
for (let i = braceStart; i < source.length; i++) {
|
||||
if (source[i] === "{") depth++;
|
||||
else if (source[i] === "}") {
|
||||
depth--;
|
||||
if (depth === 0) return source.slice(braceStart, i + 1);
|
||||
}
|
||||
}
|
||||
throw new Error(`unbalanced braces reading body of "${fnSignature}"`);
|
||||
}
|
||||
|
||||
describe("AI assistant clear/switch cancels an in-flight request (issue #5941, PR #6332)", () => {
|
||||
it("abandonInFlightRequest() invalidates the generation guard before/regardless of the backend RPC", () => {
|
||||
const body = bodyOf("function abandonInFlightRequest(alreadyCancelledSessionId?: string)");
|
||||
expect(body).toContain("aiGenerationGuard.invalidate();");
|
||||
expect(body).toContain("isGenerating.value = false;");
|
||||
expect(body).toContain('currentSessionId.value = "";');
|
||||
// Invalidation (and resetting isGenerating/currentSessionId/delta buffers) must
|
||||
// happen unconditionally and before the best-effort backend RPC, so a send()
|
||||
// that hasn't registered a session id yet is still cut off from shared state.
|
||||
expect(body.indexOf("aiGenerationGuard.invalidate();")).toBeLessThan(body.indexOf("if (sessionId && sessionId !== alreadyCancelledSessionId) {"));
|
||||
});
|
||||
|
||||
it("clearMessages() abandons a stuck stream before wiping history", () => {
|
||||
const body = bodyOf("function clearMessages()");
|
||||
expect(body).toContain("if (isGenerating.value) abandonInFlightRequest();");
|
||||
// The abandon guard must run before the history is actually wiped, not after.
|
||||
expect(body.indexOf("if (isGenerating.value) abandonInFlightRequest();")).toBeLessThan(body.indexOf("messages.value = []"));
|
||||
});
|
||||
|
||||
it("selectConversation() abandons a stuck stream before switching conversations", () => {
|
||||
const body = bodyOf("function selectConversation(conv: AiConversation)");
|
||||
expect(body).toContain("if (isGenerating.value) abandonInFlightRequest();");
|
||||
expect(body.indexOf("if (isGenerating.value) abandonInFlightRequest();")).toBeLessThan(body.indexOf("conversationId.value = conv.id;"));
|
||||
});
|
||||
|
||||
it("cancelStream() delegates live state readers to the tested stop lifecycle", () => {
|
||||
const body = bodyOf("async function cancelStream()");
|
||||
expect(body).toContain("await stopAiGenerationWithFallback({");
|
||||
expect(body).toContain("currentGeneration: () => aiGenerationGuard.peek()");
|
||||
expect(body).toContain("currentAssistantMessageIndex: () => currentAssistantMessageIndex");
|
||||
expect(body).toContain("abandon: (sessionId) => abandonInFlightRequest(sessionId)");
|
||||
expect(body).toContain("persistConversation,");
|
||||
});
|
||||
|
||||
it("startNewChat() clears through the guarded path and deleteConversation() uses the tested delete lifecycle", () => {
|
||||
expect(bodyOf("function startNewChat()")).toContain("clearMessages();");
|
||||
const deleteBody = bodyOf("async function deleteConversation(id: string)");
|
||||
expect(deleteBody).toContain("await deleteConversationWithCancellation({");
|
||||
expect(deleteBody).toContain("abandon: () => abandonInFlightRequest()");
|
||||
expect(deleteBody).toContain("deletePersisted: () => deleteAiConversation(id).catch(() => {})");
|
||||
expect(deleteBody).toContain("if (conversationId.value === id) clearMessages();");
|
||||
});
|
||||
|
||||
it("send() claims a generation id right after setting isGenerating and re-checks it after the first await", () => {
|
||||
const sendBody = bodyOf("async function send()");
|
||||
const claimIdx = sendBody.indexOf("const myGeneration = aiGenerationGuard.begin();");
|
||||
expect(claimIdx).toBeGreaterThanOrEqual(0);
|
||||
expect(sendBody.indexOf("isGenerating.value = true;")).toBeLessThan(claimIdx);
|
||||
// Must re-validate after the first await point (ensureLoaded()) before doing
|
||||
// anything that touches messages/mentions belonging to whatever conversation
|
||||
// is current by the time it resolves.
|
||||
const ensureLoadedAwaitIdx = sendBody.indexOf("await promptTemplateStore.ensureLoaded()");
|
||||
const recheckIdx = sendBody.indexOf("if (!aiGenerationGuard.isCurrent(myGeneration))", ensureLoadedAwaitIdx);
|
||||
expect(ensureLoadedAwaitIdx).toBeGreaterThan(claimIdx);
|
||||
expect(recheckIdx).toBeGreaterThan(ensureLoadedAwaitIdx);
|
||||
});
|
||||
|
||||
it("send() clears the pending write-SQL grant when superseded before it's consumed", () => {
|
||||
// Reviewed on PR #6332: at this point in send(), the write-SQL grant
|
||||
// (allowWriteSqlForNextRun/confirmedWriteSqlText/...) hasn't been read/reset
|
||||
// yet — that happens later, right before runAgentStream(). A bare `return`
|
||||
// here would leave a previously-confirmed write grant sitting in the
|
||||
// module-scope vars, live to be replayed against whatever unrelated send()
|
||||
// the next conversation issues.
|
||||
const sendBody = bodyOf("async function send()");
|
||||
const ensureLoadedAwaitIdx = sendBody.indexOf("await promptTemplateStore.ensureLoaded()");
|
||||
const grantConsumedIdx = sendBody.indexOf("const allowWriteSql = requestedMode", ensureLoadedAwaitIdx);
|
||||
const recheckIdx = sendBody.indexOf("if (!aiGenerationGuard.isCurrent(myGeneration)) {", ensureLoadedAwaitIdx);
|
||||
expect(recheckIdx).toBeGreaterThan(ensureLoadedAwaitIdx);
|
||||
expect(recheckIdx).toBeLessThan(grantConsumedIdx);
|
||||
const recheckEnd = sendBody.indexOf("}", recheckIdx);
|
||||
const recheckBlock = sendBody.slice(recheckIdx, recheckEnd);
|
||||
expect(recheckBlock).toContain("clearPendingWriteGrant();");
|
||||
expect(recheckBlock).toContain("return;");
|
||||
});
|
||||
|
||||
it("send()'s agent-event callback bails immediately if superseded", () => {
|
||||
const sendBody = bodyOf("async function send()");
|
||||
const callbackStart = sendBody.indexOf("(event: AgentEvent) => {");
|
||||
expect(callbackStart).toBeGreaterThanOrEqual(0);
|
||||
const guardIdx = sendBody.indexOf("if (!aiGenerationGuard.isCurrent(myGeneration)) return;", callbackStart);
|
||||
const pushIdx = sendBody.indexOf("agentEvents.push(event);", callbackStart);
|
||||
expect(guardIdx).toBeGreaterThan(callbackStart);
|
||||
expect(guardIdx).toBeLessThan(pushIdx);
|
||||
});
|
||||
|
||||
it("send()'s catch block guards the assistant message lookup behind the generation check", () => {
|
||||
const sendBody = bodyOf("async function send()");
|
||||
const start = sendBody.indexOf("} catch (e: unknown) {");
|
||||
expect(start, 'expected to find a "} catch (e: unknown) {" block inside send()').toBeGreaterThanOrEqual(0);
|
||||
const end = sendBody.indexOf("} finally {", start);
|
||||
const catchBody = sendBody.slice(start, end);
|
||||
// Must not index messages.value[assistantIdx] unguarded — if clearMessages()/
|
||||
// selectConversation() already replaced the array (or invalidated the
|
||||
// generation) while this request was still in flight, an unguarded write would
|
||||
// either throw (silently eating the real error) or corrupt a different
|
||||
// conversation's transcript.
|
||||
expect(catchBody).not.toContain("messages.value[assistantIdx].content =");
|
||||
expect(catchBody).toContain("aiGenerationGuard.isCurrent(myGeneration)");
|
||||
expect(catchBody).toContain("const msg = messages.value[assistantIdx];");
|
||||
expect(catchBody).toContain("if (msg) msg.content =");
|
||||
});
|
||||
|
||||
it("send()'s finally block only mutates shared state (isGenerating, currentSessionId) when still current", () => {
|
||||
const sendBody = bodyOf("async function send()");
|
||||
const start = sendBody.indexOf("} finally {");
|
||||
expect(start).toBeGreaterThanOrEqual(0);
|
||||
const finallyBody = sendBody.slice(start);
|
||||
const guardIdx = finallyBody.indexOf("if (aiGenerationGuard.isCurrent(myGeneration)) {");
|
||||
expect(guardIdx).toBeGreaterThanOrEqual(0);
|
||||
const isGeneratingIdx = finallyBody.indexOf("isGenerating.value = false;");
|
||||
const sessionResetIdx = finallyBody.indexOf('currentSessionId.value = "";');
|
||||
expect(isGeneratingIdx).toBeGreaterThan(guardIdx);
|
||||
expect(sessionResetIdx).toBeGreaterThan(guardIdx);
|
||||
});
|
||||
|
||||
// The three gaps below were called out on review of PR #6332: the generation guard
|
||||
// stopped *writes* from a superseded generation, but didn't stop the generation from
|
||||
// starting a stream, leaking pending-compaction state, or from surviving unmount.
|
||||
it("send() rechecks the generation after EACH context-preparation await, before ever starting the stream", () => {
|
||||
const sendBody = bodyOf("async function send()");
|
||||
const sqlFilesIdx = sendBody.indexOf("await loadReferencedSqlFiles(");
|
||||
const firstRecheckIdx = sendBody.indexOf("if (!aiGenerationGuard.isCurrent(myGeneration)) return;", sqlFilesIdx);
|
||||
const contextIdx = sendBody.indexOf("const context = await buildAiContext(");
|
||||
const secondRecheckIdx = sendBody.indexOf("if (!aiGenerationGuard.isCurrent(myGeneration)) return;", contextIdx);
|
||||
const runIdx = sendBody.indexOf("await runAgentStream(", contextIdx);
|
||||
// A clear/switch/unmount firing during EITHER await invalidates the generation
|
||||
// but doesn't touch the backend (no session registered yet) — without a
|
||||
// recheck immediately after each one, send() would resume into further wasted
|
||||
// work (buildAiContext() can do real backend/schema work) and eventually call
|
||||
// runAgentStream() anyway, starting a request nobody can reach anymore.
|
||||
expect(sqlFilesIdx).toBeGreaterThanOrEqual(0);
|
||||
expect(firstRecheckIdx).toBeGreaterThan(sqlFilesIdx);
|
||||
expect(firstRecheckIdx).toBeLessThan(contextIdx);
|
||||
expect(secondRecheckIdx).toBeGreaterThan(contextIdx);
|
||||
expect(runIdx).toBeGreaterThan(secondRecheckIdx);
|
||||
});
|
||||
|
||||
it("send() tracks the current assistant message index alongside currentSessionId, for cancelStream() to finalize", () => {
|
||||
// cancelStream()'s forced-abandon path needs to know which message in
|
||||
// messages.value belongs to this generation so it can finalize it (see the
|
||||
// "finalizes the stuck placeholder message" test above). Must be set at the
|
||||
// same point currentSessionId is, and cleared alongside it too.
|
||||
const sendBody = bodyOf("async function send()");
|
||||
const pushIdx = sendBody.indexOf('messages.value.push({ role: "assistant"');
|
||||
const assistantIdxIdx = sendBody.indexOf("const assistantIdx = messages.value.length - 1;", pushIdx);
|
||||
const trackIdx = sendBody.indexOf("currentAssistantMessageIndex = assistantIdx;", assistantIdxIdx);
|
||||
const sessionSetIdx = sendBody.indexOf("currentSessionId.value = sessionId;", trackIdx);
|
||||
expect(pushIdx).toBeGreaterThanOrEqual(0);
|
||||
expect(assistantIdxIdx).toBeGreaterThan(pushIdx);
|
||||
expect(trackIdx).toBeGreaterThan(assistantIdxIdx);
|
||||
expect(sessionSetIdx).toBeGreaterThan(trackIdx);
|
||||
|
||||
const start = sendBody.indexOf("} finally {");
|
||||
const finallyBody = sendBody.slice(start);
|
||||
const sessionResetIdx = finallyBody.indexOf('currentSessionId.value = "";');
|
||||
const indexResetIdx = finallyBody.indexOf("currentAssistantMessageIndex = -1;");
|
||||
expect(sessionResetIdx).toBeGreaterThanOrEqual(0);
|
||||
expect(indexResetIdx).toBeGreaterThan(sessionResetIdx);
|
||||
});
|
||||
|
||||
it("abandonInFlightRequest() skips the cancel RPC for a session the caller already cancelled", () => {
|
||||
// Reviewed on PR #6332 (efficiency): cancelStream() awaits aiCancelStream()
|
||||
// itself before deciding to force-abandon, so re-firing it unconditionally
|
||||
// here would double-RPC the same session. Must still fire for a session
|
||||
// registered AFTER the caller's own RPC attempt (i.e. not simply skipped
|
||||
// whenever a caller happened to pass an argument).
|
||||
const body = bodyOf("function abandonInFlightRequest(alreadyCancelledSessionId?: string)");
|
||||
expect(body).toContain("if (sessionId && sessionId !== alreadyCancelledSessionId) {");
|
||||
expect(body).not.toContain("if (sessionId) {\n aiCancelStream(sessionId).catch(() => {});\n }");
|
||||
});
|
||||
|
||||
it("resetPendingRequestState() resets all per-request transient state, and abandonInFlightRequest() uses it", () => {
|
||||
// Reviewed on PR #6332: this reset logic used to be duplicated inline across
|
||||
// send()'s finally, abandonInFlightRequest(), and flushAssistantDeltas() —
|
||||
// duplication that already let the pendingCompaction reset get missed once.
|
||||
// Consolidating it into one function means a newly-added piece of per-request
|
||||
// state only needs to be added here to be covered by the abandon path.
|
||||
const resetBody = bodyOf("function resetPendingRequestState()");
|
||||
expect(resetBody).toContain("cancelAnimationFrame(assistantDeltaFrame);");
|
||||
expect(resetBody).toContain("assistantDeltaFrame = null;");
|
||||
expect(resetBody).toContain('pendingAssistantDelta = "";');
|
||||
expect(resetBody).toContain('pendingAssistantReasoning = "";');
|
||||
expect(resetBody).toContain("pendingAssistantIndex = -1;");
|
||||
// A context_compacted event on the abandoned generation may have already set
|
||||
// pendingCompaction before invalidation took effect. It's request-owned, not
|
||||
// conversation-owned, so it must be reset here — otherwise the next send() (in
|
||||
// a brand-new, just-cleared conversation) would splice the OLD conversation's
|
||||
// compaction summary into the NEW conversation's transcript in its finally
|
||||
// block.
|
||||
expect(resetBody).toContain("pendingCompaction.value = null;");
|
||||
|
||||
const abandonBody = bodyOf("function abandonInFlightRequest(alreadyCancelledSessionId?: string)");
|
||||
expect(abandonBody).toContain("resetPendingRequestState();");
|
||||
});
|
||||
|
||||
it("onUnmounted() invalidates the generation instead of only firing the best-effort cancel RPC", () => {
|
||||
const body = bodyOf("onUnmounted(() => {");
|
||||
// Plain cancelStream() leaves the generation current: a request still mid-await
|
||||
// when the component unmounts would resume, call runAgentStream(), and its event
|
||||
// callback/catch/finally would keep writing into refs this now-unmounted
|
||||
// instance's closures still hold. Must go through the same invalidating path as
|
||||
// clearMessages()/selectConversation().
|
||||
expect(body).toContain("if (isGenerating.value) abandonInFlightRequest();");
|
||||
expect(body).not.toContain("cancelStream();");
|
||||
});
|
||||
});
|
||||
@@ -2382,6 +2382,7 @@ export default {
|
||||
clear: "Clear Chat",
|
||||
welcome: "Tell me what you'd like to query, and I'll write the SQL",
|
||||
requestFailed: "AI request failed",
|
||||
requestCancelled: "Request cancelled",
|
||||
newChat: "New Chat",
|
||||
thinking: "Thinking...",
|
||||
reasoningProcess: "Thinking process",
|
||||
|
||||
@@ -2421,6 +2421,7 @@ export default withEnglishFallback({
|
||||
clear: "Limpiar chat",
|
||||
welcome: "Dime qué quieres consultar y escribiré el SQL",
|
||||
requestFailed: "La solicitud de AI falló",
|
||||
requestCancelled: "Solicitud de AI cancelada",
|
||||
newChat: "Nuevo chat",
|
||||
thinking: "Procesando...",
|
||||
reasoningProcess: "Proceso de razonamiento",
|
||||
|
||||
@@ -2315,6 +2315,7 @@ export default withEnglishFallback({
|
||||
clear: "Cancella Chat",
|
||||
welcome: "Dimmi cosa vorresti interrogare e io scriverò il codice SQL",
|
||||
requestFailed: "Richiesta AI non riuscita",
|
||||
requestCancelled: "Richiesta AI annullata",
|
||||
newChat: "Nuova Chat",
|
||||
thinking: "Riflessione...",
|
||||
reasoningProcess: "Processo di riflessione",
|
||||
|
||||
@@ -2455,6 +2455,7 @@ export default withEnglishFallback({
|
||||
clear: "チャットをクリア",
|
||||
welcome: "クエリしたい内容を教えてください。SQLを作成します",
|
||||
requestFailed: "AIリクエストに失敗しました",
|
||||
requestCancelled: "AIリクエストがキャンセルされました",
|
||||
newChat: "新しいチャット",
|
||||
thinking: "思考中...",
|
||||
reasoningProcess: "思考プロセス",
|
||||
|
||||
@@ -2232,6 +2232,7 @@ export default withEnglishFallback({
|
||||
clear: "채팅 지우기",
|
||||
welcome: "쿼리할 내용을 말씀해 주세요. SQL을 작성해 드릴게요",
|
||||
requestFailed: "AI 요청 실패",
|
||||
requestCancelled: "AI 요청이 취소되었습니다",
|
||||
newChat: "새 채팅",
|
||||
thinking: "생각하는 중...",
|
||||
reasoningProcess: "사고 과정",
|
||||
|
||||
@@ -2423,6 +2423,7 @@ export default withEnglishFallback({
|
||||
clear: "Limpar Conversa",
|
||||
welcome: "Diga o que você gostaria de consultar e eu escreverei o SQL",
|
||||
requestFailed: "Falha na solicitação de AI",
|
||||
requestCancelled: "Solicitação de AI cancelada",
|
||||
newChat: "Nova Conversa",
|
||||
thinking: "Pensando...",
|
||||
reasoningProcess: "Processo de raciocínio",
|
||||
|
||||
@@ -2306,6 +2306,7 @@ export default withEnglishFallback({
|
||||
clear: "清空对话",
|
||||
welcome: "聊聊你想查什么，我来写 SQL",
|
||||
requestFailed: "AI 请求失败",
|
||||
requestCancelled: "请求已取消",
|
||||
newChat: "新对话",
|
||||
thinking: "思考中...",
|
||||
reasoningProcess: "思考过程",
|
||||
|
||||
@@ -2316,6 +2316,7 @@ export default withEnglishFallback({
|
||||
clear: "清空對話",
|
||||
welcome: "告訴我你想查詢什麼，我會寫出 SQL",
|
||||
requestFailed: "AI 請求失敗",
|
||||
requestCancelled: "AI 請求已取消",
|
||||
newChat: "新對話",
|
||||
thinking: "思考中……",
|
||||
reasoningProcess: "思考過程",
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { deleteConversationWithCancellation, stopAiGenerationWithFallback } from "@/lib/ai/aiConversationLifecycle";
|
||||
|
||||
function deferred() {
|
||||
let resolve!: () => void;
|
||||
const promise = new Promise<void>((done) => {
|
||||
resolve = done;
|
||||
});
|
||||
return { promise, resolve };
|
||||
}
|
||||
|
||||
describe("AI conversation lifecycle", () => {
|
||||
it("abandons the active generation before deleting its conversation", async () => {
|
||||
const deletion = deferred();
|
||||
const events: string[] = [];
|
||||
|
||||
const pending = deleteConversationWithCancellation({
|
||||
id: "conversation-1",
|
||||
currentConversationId: () => "conversation-1",
|
||||
isGenerating: () => true,
|
||||
abandon: () => events.push("abandon"),
|
||||
deletePersisted: () => {
|
||||
events.push("delete");
|
||||
return deletion.promise;
|
||||
},
|
||||
afterDelete: () => events.push("after-delete"),
|
||||
});
|
||||
|
||||
expect(events).toEqual(["abandon", "delete"]);
|
||||
deletion.resolve();
|
||||
await pending;
|
||||
expect(events).toEqual(["abandon", "delete", "after-delete"]);
|
||||
});
|
||||
|
||||
it("uses the assistant placeholder created after Stop was clicked", async () => {
|
||||
const wait = deferred();
|
||||
const messages = [{ content: "question" }, { content: "", isThinking: true }];
|
||||
let assistantMessageIndex = -1;
|
||||
const abandon = vi.fn();
|
||||
const persistConversation = vi.fn().mockResolvedValue(undefined);
|
||||
|
||||
const pending = stopAiGenerationWithFallback({
|
||||
isGenerating: () => true,
|
||||
currentGeneration: () => 7,
|
||||
isGenerationCurrent: (generation) => generation === 7,
|
||||
currentSessionId: () => "",
|
||||
cancelSession: vi.fn().mockResolvedValue(undefined),
|
||||
waitForGenerationToClear: () => wait.promise,
|
||||
flushPending: vi.fn(),
|
||||
currentAssistantMessageIndex: () => assistantMessageIndex,
|
||||
messageAt: (index) => messages[index],
|
||||
cancelledMessage: () => "Request cancelled",
|
||||
abandon,
|
||||
persistConversation,
|
||||
});
|
||||
|
||||
assistantMessageIndex = 1;
|
||||
wait.resolve();
|
||||
await pending;
|
||||
|
||||
expect(messages[1]).toEqual({ content: "Request cancelled", isThinking: false });
|
||||
expect(abandon).toHaveBeenCalledWith("");
|
||||
expect(persistConversation).toHaveBeenCalledOnce();
|
||||
});
|
||||
|
||||
it("persists the finalized partial response after forced abandon", async () => {
|
||||
const message = { content: "partial", isThinking: true };
|
||||
const events: string[] = [];
|
||||
let persistedMessage: typeof message | undefined;
|
||||
|
||||
await stopAiGenerationWithFallback({
|
||||
isGenerating: () => true,
|
||||
currentGeneration: () => 11,
|
||||
isGenerationCurrent: (generation) => generation === 11,
|
||||
currentSessionId: () => "session-1",
|
||||
cancelSession: async () => {},
|
||||
waitForGenerationToClear: async () => {},
|
||||
flushPending: () => events.push("flush"),
|
||||
currentAssistantMessageIndex: () => 0,
|
||||
messageAt: () => message,
|
||||
cancelledMessage: () => "Request cancelled",
|
||||
abandon: () => events.push("abandon"),
|
||||
persistConversation: async () => {
|
||||
events.push("persist");
|
||||
persistedMessage = { ...message };
|
||||
},
|
||||
});
|
||||
|
||||
expect(events).toEqual(["flush", "abandon", "persist"]);
|
||||
expect(persistedMessage).toEqual({ content: "partial", isThinking: false });
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,110 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { AiGenerationGuard } from "@/lib/ai/aiGenerationGuard";
|
||||
|
||||
// Race coverage requested on https://github.com/t8y2/dbx/pull/6332: clearing/switching
|
||||
// the AI chat while a request is in flight must isolate that request's async callbacks
|
||||
// from whatever generation is active by the time they run, independent of whether the
|
||||
// backend cancel RPC actually reached a registered stream id.
|
||||
|
||||
describe("AiGenerationGuard", () => {
|
||||
it("begin() returns increasing ids and starts current", () => {
|
||||
const guard = new AiGenerationGuard();
|
||||
const g1 = guard.begin();
|
||||
expect(guard.isCurrent(g1)).toBe(true);
|
||||
const g2 = guard.begin();
|
||||
expect(g2).toBeGreaterThan(g1);
|
||||
expect(guard.isCurrent(g1)).toBe(false);
|
||||
expect(guard.isCurrent(g2)).toBe(true);
|
||||
});
|
||||
|
||||
it("invalidate() supersedes the active generation even with no new send()", () => {
|
||||
const guard = new AiGenerationGuard();
|
||||
const g1 = guard.begin();
|
||||
guard.invalidate();
|
||||
expect(guard.isCurrent(g1)).toBe(false);
|
||||
});
|
||||
|
||||
// peek() lets a caller that isn't itself a send() invocation (the Stop button —
|
||||
// see AiAssistant.vue's cancelStream()) snapshot the active generation without
|
||||
// starting a new one, then later check isCurrent() against that snapshot. This
|
||||
// is what lets Stop detect "did the request I clicked on get superseded" even
|
||||
// when there's no session id registered yet to compare against instead
|
||||
// (reviewed on PR #6332).
|
||||
it("peek() snapshots the current generation without advancing it", () => {
|
||||
const guard = new AiGenerationGuard();
|
||||
const g1 = guard.begin();
|
||||
const snapshot = guard.peek();
|
||||
expect(snapshot).toBe(g1);
|
||||
expect(guard.isCurrent(snapshot)).toBe(true);
|
||||
// peek() itself must not consume/advance the generation the way begin() does.
|
||||
expect(guard.peek()).toBe(snapshot);
|
||||
});
|
||||
|
||||
it("peek() reflects invalidation the same way isCurrent() does", () => {
|
||||
const guard = new AiGenerationGuard();
|
||||
const clickedGeneration = guard.peek();
|
||||
expect(guard.isCurrent(clickedGeneration)).toBe(true);
|
||||
// A newer send() starting (begin()) or an abandon (invalidate()) after the
|
||||
// click must both make the snapshot stale.
|
||||
guard.begin();
|
||||
expect(guard.isCurrent(clickedGeneration)).toBe(false);
|
||||
});
|
||||
|
||||
// Race 1: clear/switch fires before send() has registered a stream id with the
|
||||
// backend (cancelStream() would be a no-op in this window). The guard must still
|
||||
// mark the generation stale immediately and synchronously.
|
||||
it("invalidating before stream registration still supersedes the pending send()", async () => {
|
||||
const guard = new AiGenerationGuard();
|
||||
let registeredWhileCurrent: boolean | null = null;
|
||||
|
||||
async function send() {
|
||||
const gen = guard.begin();
|
||||
// Simulate the async work send() does before it registers a session id
|
||||
// with the backend (e.g. awaiting promptTemplateStore.ensureLoaded()).
|
||||
await Promise.resolve();
|
||||
// The user cleared the chat during that await, before registration.
|
||||
registeredWhileCurrent = guard.isCurrent(gen);
|
||||
}
|
||||
|
||||
const pending = send();
|
||||
// Clear fires synchronously, before send()'s microtask queue drains —
|
||||
// i.e. before any session id could have been registered.
|
||||
guard.invalidate();
|
||||
await pending;
|
||||
|
||||
expect(registeredWhileCurrent).toBe(false);
|
||||
});
|
||||
|
||||
// Race 2: cancel, then immediately start a new send() — the old request's promise
|
||||
// resolves later. The old generation's completion handlers must not touch state
|
||||
// that now belongs to the new generation.
|
||||
it("an old generation resolving after a newer one started must not report current", async () => {
|
||||
const guard = new AiGenerationGuard();
|
||||
const applied: string[] = [];
|
||||
|
||||
function startSend(label: string, resolveAfter: Promise<void>) {
|
||||
const gen = guard.begin();
|
||||
return resolveAfter.then(() => {
|
||||
if (guard.isCurrent(gen)) applied.push(label);
|
||||
});
|
||||
}
|
||||
|
||||
let releaseOld: () => void;
|
||||
const oldFinishes = new Promise<void>((resolve) => {
|
||||
releaseOld = resolve;
|
||||
});
|
||||
const oldSend = startSend("old", oldFinishes);
|
||||
|
||||
// User cancels (clear/switch) and immediately fires a new request before the
|
||||
// old one has resolved.
|
||||
guard.invalidate();
|
||||
const newSend = startSend("new", Promise.resolve());
|
||||
await newSend;
|
||||
|
||||
// Old request's backend call finally settles well after the new one started.
|
||||
releaseOld!();
|
||||
await oldSend;
|
||||
|
||||
expect(applied).toEqual(["new"]);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,55 @@
|
||||
export interface AiAssistantLifecycleMessage {
|
||||
content: string;
|
||||
isThinking?: boolean;
|
||||
}
|
||||
|
||||
interface StopAiGenerationOptions<TGeneration, TMessage extends AiAssistantLifecycleMessage> {
|
||||
isGenerating: () => boolean;
|
||||
currentGeneration: () => TGeneration;
|
||||
isGenerationCurrent: (generation: TGeneration) => boolean;
|
||||
currentSessionId: () => string;
|
||||
cancelSession: (sessionId: string) => Promise<void>;
|
||||
waitForGenerationToClear: () => Promise<void>;
|
||||
flushPending: () => void;
|
||||
currentAssistantMessageIndex: () => number;
|
||||
messageAt: (index: number) => TMessage | undefined;
|
||||
cancelledMessage: () => string;
|
||||
abandon: (alreadyCancelledSessionId: string) => void;
|
||||
persistConversation: () => Promise<void>;
|
||||
}
|
||||
|
||||
export async function stopAiGenerationWithFallback<TGeneration, TMessage extends AiAssistantLifecycleMessage>(options: StopAiGenerationOptions<TGeneration, TMessage>): Promise<boolean> {
|
||||
if (!options.isGenerating()) return false;
|
||||
|
||||
const generation = options.currentGeneration();
|
||||
const sessionId = options.currentSessionId();
|
||||
if (sessionId) await options.cancelSession(sessionId).catch(() => {});
|
||||
|
||||
await options.waitForGenerationToClear();
|
||||
if (!options.isGenerating() || !options.isGenerationCurrent(generation)) return false;
|
||||
|
||||
options.flushPending();
|
||||
const message = options.messageAt(options.currentAssistantMessageIndex());
|
||||
if (message) {
|
||||
message.isThinking = false;
|
||||
if (!message.content) message.content = options.cancelledMessage();
|
||||
}
|
||||
options.abandon(sessionId);
|
||||
await options.persistConversation();
|
||||
return true;
|
||||
}
|
||||
|
||||
interface DeleteConversationOptions {
|
||||
id: string;
|
||||
currentConversationId: () => string;
|
||||
isGenerating: () => boolean;
|
||||
abandon: () => void;
|
||||
deletePersisted: () => Promise<void>;
|
||||
afterDelete: () => void;
|
||||
}
|
||||
|
||||
export async function deleteConversationWithCancellation(options: DeleteConversationOptions): Promise<void> {
|
||||
if (options.currentConversationId() === options.id && options.isGenerating()) options.abandon();
|
||||
await options.deletePersisted();
|
||||
options.afterDelete();
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
/**
|
||||
* Tracks which in-flight `send()` invocation ("generation") is still allowed to
|
||||
* write into the AI assistant's shared reactive state (`messages`, `isGenerating`,
|
||||
* `currentSessionId`, the streaming delta buffers).
|
||||
*
|
||||
* `messages`/`isGenerating` are plain component refs, not per-conversation state.
|
||||
* Clearing the chat, switching to a different saved conversation, starting a new
|
||||
* one, or unmounting the component replaces/abandons `messages.value` and must
|
||||
* invalidate whatever `send()` call is still running — otherwise its async event
|
||||
* callbacks, `catch`, or `finally` can keep writing into the array that now
|
||||
* belongs to a different conversation (or to an unmounted instance), or clear
|
||||
* `isGenerating` for a request that has nothing to do with them.
|
||||
*
|
||||
* The backend cancel RPC is a best-effort companion to this guard, not a
|
||||
* substitute for it: it depends on the stream having already registered a
|
||||
* session id, which happens partway through `send()`. If clear/switch fires
|
||||
* before that registration, the backend call is a no-op — but `invalidate()`
|
||||
* still takes effect immediately and synchronously, so the superseded
|
||||
* generation's callbacks stop touching shared state regardless of whether the
|
||||
* backend was actually told to stop.
|
||||
*/
|
||||
export class AiGenerationGuard {
|
||||
private current = 0;
|
||||
|
||||
/** Starts a new generation and returns its id. Call once per `send()` invocation. */
|
||||
begin(): number {
|
||||
this.current += 1;
|
||||
return this.current;
|
||||
}
|
||||
|
||||
/** Invalidates whatever generation is currently active (clear / switch / new chat / explicit cancel). */
|
||||
invalidate(): void {
|
||||
this.current += 1;
|
||||
}
|
||||
|
||||
/** True when `generation` is still the active one and may write shared state. */
|
||||
isCurrent(generation: number): boolean {
|
||||
return generation === this.current;
|
||||
}
|
||||
|
||||
/**
|
||||
* Snapshots the active generation without starting a new one. Lets a caller that
|
||||
* isn't itself a `send()` invocation (e.g. the Stop button) later ask "is the
|
||||
* request I observed at snapshot time still the one running?" via `isCurrent()`
|
||||
* — this is what lets that check work even when there's no session id yet to
|
||||
* compare against instead.
|
||||
*/
|
||||
peek(): number {
|
||||
return this.current;
|
||||
}
|
||||
}
|
||||
+5
-1
@@ -60,7 +60,11 @@ test("assistant messages keep metadata aligned with the bubble and wrap long tex
|
||||
});
|
||||
|
||||
test("AI request failures use localized backend diagnostics", () => {
|
||||
assert.match(source, /messages\.value\[assistantIdx\]\.content = `\$\{t\("ai\.requestFailed"\)\}\\n\\n\$\{translateBackendError\(t, message\)\}`/);
|
||||
// The lookup is guarded (matching the `finally` block right below it) so that if
|
||||
// clearMessages()/selectConversation() already replaced messages.value while this
|
||||
// request was still in flight, writing the failure text doesn't throw and silently
|
||||
// swallow the real error. See issue #5941.
|
||||
assert.match(source, /const msg = messages\.value\[assistantIdx\];\s*\n\s*if \(msg\) msg\.content = `\$\{t\("ai\.requestFailed"\)\}\\n\\n\$\{translateBackendError\(t, message\)\}`;/);
|
||||
});
|
||||
|
||||
test("AI analysis export keeps the connection that produced each assistant response", () => {
|
||||
|
||||
Reference in New Issue
Block a user