fix(ai): cancel active streams before conversation changes

apps/desktop/src/lib/__tests__/ai/aiConversationLifecycle.spec.ts

@@ -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 });
+  });
+});

apps/desktop/src/lib/__tests__/ai/aiGenerationGuard.spec.ts

@@ -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"]);
+  });
+});

apps/desktop/src/lib/ai/aiConversationLifecycle.ts

@@ -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();
+}

apps/desktop/src/lib/ai/aiGenerationGuard.ts

@@ -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;
+  }
+}