fix(ai): cancel active streams before conversation changes
This commit is contained in:
1 parent
e6fa850a31
commit
3e6e43c57c
15 files changed
+805
-53
No files matched your search
@@ -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();");
|
||||
});
|
||||
});
|
||||
Reference in new issue
Block a user