fix(ai-chat): open chat window before resolving the bound chat (#191)

Address PR #209 review.

- use-open-ai-chat.ts: call setWindowOpen(true) before awaiting
  getBoundChat so the header button feels instant on slow connections;
  the chat switch (setActiveChatId/setDraft/setSelectedRoleId) is applied
  after the round-trip resolves. Also drop the redundant no-op
  setWindowOpen(true) in the already-open branch (bare early return).
- CHANGELOG.md: document the header AI-chat button auto-opening the
  latest chat bound to the current document under [Unreleased]/Added.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

CHANGELOG.md

@@ -72,6 +72,13 @@ per-workspace rolling-day token budget.
 - **Footnote multi-backlinks.** A footnote referenced more than once now shows a
   back-link per reference (↩ a b c …), each scrolling to its own occurrence, like
   Pandoc/Wikipedia; a single-reference footnote keeps the plain ↩. (#168)
+- **AI chat: header button auto-opens the chat bound to the current document.**
+  Clicking the AI-chat button in the header while viewing a page now reopens the
+  latest chat tied to that document instead of whatever chat was last active,
+  reusing the existing `ai_chats.page_id` provenance (no migration). The newest
+  chat you created on the page wins; with no bound chat — or off a page, or if
+  the lookup fails — it falls soft to a fresh chat and keeps the current
+  selection otherwise. (#191)
 
 ### Changed
 

apps/client/src/components/layouts/global/app-header.tsx

@@ -10,12 +10,12 @@ import classes from "./app-header.module.css";
 import { BrandLogo } from "@/components/ui/brand-logo";
 import TopMenu from "@/components/layouts/global/top-menu.tsx";
 import { Link } from "react-router-dom";
-import { useAtom, useSetAtom } from "jotai";
+import { useAtom } from "jotai";
 import {
   desktopSidebarAtom,
   mobileSidebarAtom,
 } from "@/components/layouts/global/hooks/atoms/sidebar-atom.ts";
-import { aiChatWindowOpenAtom } from "@/features/ai-chat/atoms/ai-chat-atom.ts";
+import { useOpenAiChatForCurrentPage } from "@/features/ai-chat/hooks/use-open-ai-chat.ts";
 import { workspaceAtom } from "@/features/user/atoms/current-user-atom.ts";
 import { useToggleSidebar } from "@/components/layouts/global/hooks/hooks/use-toggle-sidebar.ts";
 import SidebarToggle from "@/components/ui/sidebar-toggle-button.tsx";
@@ -38,7 +38,9 @@ export function AppHeader() {
   const toggleDesktop = useToggleSidebar(desktopSidebarAtom);
 
   const [workspace] = useAtom(workspaceAtom);
-  const setAiChatWindowOpen = useSetAtom(aiChatWindowOpenAtom);
+  // Opening from the header auto-opens the document's bound chat (last chat
+  // created on the current page); off a page it keeps the current selection.
+  const openAiChat = useOpenAiChatForCurrentPage();
   // AI chat entry point: only shown when the workspace enables it (A7 gate).
   const aiChatEnabled = workspace?.settings?.ai?.chat === true;
 
@@ -105,7 +107,7 @@ export function AppHeader() {
                 color="dark"
                 size="sm"
                 aria-label={t("AI chat")}
-                onClick={() => setAiChatWindowOpen((v) => !v)}
+                onClick={openAiChat}
               >
                 <IconMessage size={20} />
               </ActionIcon>

apps/client/src/features/ai-chat/hooks/use-open-ai-chat.test.tsx

@@ -0,0 +1,135 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, act } from "@testing-library/react";
+import { Provider, createStore } from "jotai";
+import type { ReactNode } from "react";
+import { useOpenAiChatForCurrentPage } from "./use-open-ai-chat";
+import {
+  activeAiChatIdAtom,
+  aiChatWindowOpenAtom,
+  aiChatDraftAtom,
+  selectedAiRoleIdAtom,
+} from "@/features/ai-chat/atoms/ai-chat-atom.ts";
+
+// useMatch is the only react-router-dom export the hook uses; drive its return
+// per test to simulate "on a page" vs "off a page".
+const useMatchMock = vi.fn();
+vi.mock("react-router-dom", () => ({
+  useMatch: () => useMatchMock(),
+}));
+
+// The bound-chat resolver is the network boundary; stub it per test.
+const getBoundChatMock = vi.fn();
+vi.mock("@/features/ai-chat/services/ai-chat-service.ts", () => ({
+  getBoundChat: (pageId: string) => getBoundChatMock(pageId),
+}));
+
+// Put the hook on a page route by default ("doc-p1" -> page id "p1"); individual
+// tests override useMatch to go off-page.
+function onPage(pageSlug = "doc-p1") {
+  useMatchMock.mockReturnValue({ params: { pageSlug } });
+}
+function offPage() {
+  useMatchMock.mockReturnValue(null);
+}
+
+// Render the hook inside an explicit jotai store so atom side effects are
+// assertable; the store is returned for setup + assertions.
+function setup(seed?: (store: ReturnType<typeof createStore>) => void) {
+  const store = createStore();
+  seed?.(store);
+  const wrapper = ({ children }: { children: ReactNode }) => (
+    <Provider store={store}>{children}</Provider>
+  );
+  const { result } = renderHook(() => useOpenAiChatForCurrentPage(), { wrapper });
+  return { store, open: () => act(() => result.current()) };
+}
+
+describe("useOpenAiChatForCurrentPage", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    onPage();
+  });
+
+  it("on a page: resolves the bound chat, selects it, and opens the window", async () => {
+    getBoundChatMock.mockResolvedValue("bound-chat-1");
+    const { store, open } = setup((s) => s.set(aiChatDraftAtom, "stale draft"));
+
+    await open();
+
+    expect(getBoundChatMock).toHaveBeenCalledWith("p1");
+    expect(store.get(activeAiChatIdAtom)).toBe("bound-chat-1");
+    expect(store.get(aiChatWindowOpenAtom)).toBe(true);
+    expect(store.get(aiChatDraftAtom)).toBe(""); // cleared on a real switch
+  });
+
+  it("on a page with no bound chat: opens a fresh chat (null)", async () => {
+    getBoundChatMock.mockResolvedValue(null);
+    const { store, open } = setup((s) => s.set(activeAiChatIdAtom, "previous"));
+
+    await open();
+
+    expect(store.get(activeAiChatIdAtom)).toBeNull();
+    expect(store.get(aiChatWindowOpenAtom)).toBe(true);
+  });
+
+  it("off a page: keeps the current selection and does NOT resolve", async () => {
+    offPage();
+    const { store, open } = setup((s) => {
+      s.set(activeAiChatIdAtom, "keep-me");
+      s.set(aiChatDraftAtom, "untouched");
+    });
+
+    await open();
+
+    expect(getBoundChatMock).not.toHaveBeenCalled();
+    expect(store.get(activeAiChatIdAtom)).toBe("keep-me");
+    expect(store.get(aiChatDraftAtom)).toBe("untouched"); // no switch -> kept
+    expect(store.get(aiChatWindowOpenAtom)).toBe(true);
+  });
+
+  it("window already open: re-click does NOT re-resolve or switch chats", async () => {
+    getBoundChatMock.mockResolvedValue("would-switch");
+    const { store, open } = setup((s) => {
+      s.set(aiChatWindowOpenAtom, true);
+      s.set(activeAiChatIdAtom, "current");
+    });
+
+    await open();
+
+    expect(getBoundChatMock).not.toHaveBeenCalled();
+    expect(store.get(activeAiChatIdAtom)).toBe("current");
+    expect(store.get(aiChatWindowOpenAtom)).toBe(true);
+  });
+
+  it("does NOT clear the draft when the resolved chat equals the current one", async () => {
+    getBoundChatMock.mockResolvedValue("same");
+    const { store, open } = setup((s) => {
+      s.set(activeAiChatIdAtom, "same");
+      s.set(aiChatDraftAtom, "in-progress");
+    });
+
+    await open();
+
+    expect(store.get(aiChatDraftAtom)).toBe("in-progress"); // no switch
+    expect(store.get(aiChatWindowOpenAtom)).toBe(true);
+  });
+
+  it("fail-soft: a resolve error opens a fresh chat (null)", async () => {
+    getBoundChatMock.mockRejectedValue(new Error("network"));
+    const { store, open } = setup((s) => s.set(activeAiChatIdAtom, "previous"));
+
+    await open();
+
+    expect(store.get(activeAiChatIdAtom)).toBeNull();
+    expect(store.get(aiChatWindowOpenAtom)).toBe(true);
+  });
+
+  it("clears the picked role on a real switch", async () => {
+    getBoundChatMock.mockResolvedValue("bound");
+    const { store, open } = setup((s) => s.set(selectedAiRoleIdAtom, "role-1"));
+
+    await open();
+
+    expect(store.get(selectedAiRoleIdAtom)).toBeNull();
+  });
+});

apps/client/src/features/ai-chat/hooks/use-open-ai-chat.ts

@@ -0,0 +1,67 @@
+import { useCallback } from "react";
+import { useAtom, useSetAtom } from "jotai";
+import { useMatch } from "react-router-dom";
+import {
+  aiChatWindowOpenAtom,
+  activeAiChatIdAtom,
+  aiChatDraftAtom,
+  selectedAiRoleIdAtom,
+} from "@/features/ai-chat/atoms/ai-chat-atom.ts";
+import { getBoundChat } from "@/features/ai-chat/services/ai-chat-service.ts";
+import { extractPageSlugId } from "@/lib";
+
+/**
+ * The generic "open the AI chat" action, WITH document binding: when invoked
+ * while viewing a page, it resolves that page's bound chat and selects it before
+ * opening — so the last chat for this document re-opens by itself. With no bound
+ * chat (or off a page) it keeps the current selection / opens a fresh chat. Used
+ * by the app-header entry point; NOT by the provenance badge (which deep-links).
+ */
+export function useOpenAiChatForCurrentPage() {
+  const [windowOpen, setWindowOpen] = useAtom(aiChatWindowOpenAtom);
+  const [activeChatId, setActiveChatId] = useAtom(activeAiChatIdAtom);
+  const setDraft = useSetAtom(aiChatDraftAtom);
+  const setSelectedRoleId = useSetAtom(selectedAiRoleIdAtom);
+
+  // Same route-match trick the window uses: read :pageSlug from the pathname.
+  // AiChatWindow lives in a pathless parent layout route, so useParams() can't
+  // see :pageSlug — match the full path against the authenticated page route.
+  const match = useMatch("/s/:spaceSlug/p/:pageSlug");
+  const pageId = extractPageSlugId(match?.params?.pageSlug);
+
+  return useCallback(async () => {
+    // Re-clicks while the window is already open (incl. minimized) must NOT
+    // re-resolve and yank the user to another chat: resolve only on a genuine
+    // closed -> open transition. (`windowOpen` is already true here, so there
+    // is nothing to set — just bail.)
+    if (windowOpen) return;
+    // Open the window FIRST so the control feels instant: the bound-chat
+    // round-trip below must never gate the window appearing, or on a slow
+    // connection the first click reads as a hung control until the POST returns.
+    setWindowOpen(true);
+    let resolved: string | null = activeChatId; // off-a-page: keep current
+    if (pageId) {
+      try {
+        resolved = await getBoundChat(pageId); // null => fresh chat
+      } catch {
+        resolved = null; // fail-soft: a fresh chat is always a safe fallback
+      }
+    }
+    // Clear the composer draft / picked role ONLY on an actual switch, so
+    // reopening the same chat does not wipe an in-progress draft. Applied after
+    // the resolve so the window is already visible while the switch settles.
+    if (resolved !== activeChatId) {
+      setActiveChatId(resolved);
+      setDraft("");
+      setSelectedRoleId(null);
+    }
+  }, [
+    windowOpen,
+    activeChatId,
+    pageId,
+    setWindowOpen,
+    setActiveChatId,
+    setDraft,
+    setSelectedRoleId,
+  ]);
+}

apps/client/src/features/ai-chat/services/ai-chat-service.ts

@@ -37,6 +37,17 @@ export async function getAiChatMessages(
   return req.data;
 }
 
+/**
+ * Resolve the chat bound to a document (the current user's most-recent chat
+ * created on that page), or null when there is none. Drives auto-open-on-page.
+ */
+export async function getBoundChat(pageId: string): Promise<string | null> {
+  const req = await api.post<{ chatId: string | null }>("/ai-chat/bound-chat", {
+    pageId,
+  });
+  return req.data.chatId;
+}
+
 /** Rename a chat. */
 export async function renameAiChat(data: {
   chatId: string;

apps/server/src/core/ai-chat/ai-chat.controller.bound-chat.spec.ts

@@ -0,0 +1,44 @@
+import { AiChatController } from './ai-chat.controller';
+import type { User, Workspace } from '@docmost/db/types/entity.types';
+
+/**
+ * Wiring spec for the #191 `POST /ai-chat/bound-chat` endpoint. It must forward
+ * the requesting user + workspace + pageId to findLatestByPage and return the
+ * matched chat's id, or `{ chatId: null }` when there is none. The repo already
+ * scopes to the caller's OWN chats, so a foreign pageId simply yields no match
+ * (null) — no extra page-access check is needed. Exercised with hand-rolled
+ * mocks, no Nest graph and no DB.
+ */
+describe('AiChatController.boundChat', () => {
+  const user = { id: 'u1' } as User;
+  const workspace = { id: 'ws1' } as Workspace;
+
+  function makeController(chat: unknown) {
+    const aiChatRepo = {
+      findLatestByPage: jest.fn().mockResolvedValue(chat),
+    };
+    const controller = new AiChatController(
+      {} as never,
+      aiChatRepo as never,
+      {} as never,
+      {} as never,
+    );
+    return { controller, aiChatRepo };
+  }
+
+  it('returns the owned chat id and scopes the lookup to user + workspace + page', async () => {
+    const { controller, aiChatRepo } = makeController({
+      id: 'c1',
+      creatorId: 'u1',
+    });
+    const res = await controller.boundChat({ pageId: 'p1' }, user, workspace);
+    expect(aiChatRepo.findLatestByPage).toHaveBeenCalledWith('u1', 'ws1', 'p1');
+    expect(res).toEqual({ chatId: 'c1' });
+  });
+
+  it('returns { chatId: null } for a page with no owned chat (incl. foreign pageId)', async () => {
+    const { controller } = makeController(undefined);
+    const res = await controller.boundChat({ pageId: 'foreign' }, user, workspace);
+    expect(res).toEqual({ chatId: null });
+  });
+});

apps/server/src/core/ai-chat/ai-chat.controller.ts

@@ -30,6 +30,7 @@ import { FileInterceptor } from '../../common/interceptors/file.interceptor';
 import { AiChatService, AiChatStreamBody } from './ai-chat.service';
 import { AiTranscriptionService } from './ai-transcription.service';
 import {
+  BoundChatDto,
   ChatIdDto,
   ExportChatDto,
   GetChatMessagesDto,
@@ -66,6 +67,28 @@ export class AiChatController {
     return this.aiChatRepo.findByCreator(user.id, workspace.id, pagination);
   }
 
+  /**
+   * Resolve the chat bound to a document for the requesting user: the most-recent
+   * non-deleted chat created on that page (ai_chats.page_id). Returns
+   * { chatId: null } when the page has no owned chat (-> a fresh chat). No page
+   * access check needed: only the caller's OWN chats are matched, so a foreign
+   * pageId reveals nothing.
+   */
+  @HttpCode(HttpStatus.OK)
+  @Post('bound-chat')
+  async boundChat(
+    @Body() dto: BoundChatDto,
+    @AuthUser() user: User,
+    @AuthWorkspace() workspace: Workspace,
+  ): Promise<{ chatId: string | null }> {
+    const chat = await this.aiChatRepo.findLatestByPage(
+      user.id,
+      workspace.id,
+      dto.pageId,
+    );
+    return { chatId: chat?.id ?? null };
+  }
+
   /** Fetch the messages of a chat (oldest first, paginated). */
   @HttpCode(HttpStatus.OK)
   @Post('messages')

apps/server/src/core/ai-chat/dto/ai-chat.dto.ts

@@ -27,6 +27,12 @@ export class GetChatMessagesDto {
   cursor?: string;
 }
 
+/** Resolve the chat bound to a document (the page's most-recent owned chat). */
+export class BoundChatDto {
+  @IsString()
+  pageId: string;
+}
+
 /** Export a chat to Markdown (#183). `lang` localizes the few fixed
  *  role/tool-action labels; defaults to English server-side. */
 export class ExportChatDto {

apps/server/src/database/repos/ai-chat/ai-chat.repo.spec.ts

@@ -0,0 +1,85 @@
+import { AiChatRepo } from './ai-chat.repo';
+import type { KyselyDB } from '../../types/kysely.types';
+
+/**
+ * Unit test for AiChatRepo.findLatestByPage — the "bound chat" resolver behind
+ * #191 (auto-open the last chat created on a document). It builds the scoping
+ * query, so we assert the EXACT predicates/ordering the spec mandates over a
+ * chainable builder mock (no live DB): user + workspace + page scope, the
+ * deletedAt filter, newest-by-createdAt with an id tiebreaker, limit 1. A
+ * live-Postgres ordering test is out of scope for this pure unit test.
+ */
+describe('AiChatRepo.findLatestByPage', () => {
+  type Recorded = {
+    table?: string;
+    wheres: Array<[string, string, unknown]>;
+    orderBys: Array<[string, string]>;
+    limit?: number;
+  };
+
+  function makeDb(result: unknown): { db: KyselyDB; rec: Recorded } {
+    const rec: Recorded = { wheres: [], orderBys: [] };
+    const builder: Record<string, unknown> = {};
+    const chain = () => builder;
+    builder.selectAll = chain;
+    builder.where = (col: string, op: string, val: unknown) => {
+      rec.wheres.push([col, op, val]);
+      return builder;
+    };
+    builder.orderBy = (col: string, dir: string) => {
+      rec.orderBys.push([col, dir]);
+      return builder;
+    };
+    builder.limit = (n: number) => {
+      rec.limit = n;
+      return builder;
+    };
+    builder.executeTakeFirst = () => Promise.resolve(result);
+    const db = {
+      selectFrom: (table: string) => {
+        rec.table = table;
+        return builder;
+      },
+    } as unknown as KyselyDB;
+    return { db, rec };
+  }
+
+  it('returns the matched chat and scopes by user + workspace + page (deletedAt null)', async () => {
+    const chat = { id: 'c1', creatorId: 'u1', workspaceId: 'ws1', pageId: 'p1' };
+    const { db, rec } = makeDb(chat);
+    const repo = new AiChatRepo(db);
+
+    const res = await repo.findLatestByPage('u1', 'ws1', 'p1');
+
+    expect(res).toBe(chat);
+    expect(rec.table).toBe('aiChats');
+    expect(rec.wheres).toEqual(
+      expect.arrayContaining([
+        ['creatorId', '=', 'u1'],
+        ['workspaceId', '=', 'ws1'],
+        ['pageId', '=', 'p1'],
+        ['deletedAt', 'is', null],
+      ]),
+    );
+  });
+
+  it('orders newest-first by createdAt then id, limit 1', async () => {
+    const { db, rec } = makeDb(undefined);
+    const repo = new AiChatRepo(db);
+
+    await repo.findLatestByPage('u1', 'ws1', 'p1');
+
+    expect(rec.orderBys).toEqual([
+      ['createdAt', 'desc'],
+      ['id', 'desc'],
+    ]);
+    expect(rec.limit).toBe(1);
+  });
+
+  it('returns undefined when the page has no owned chat', async () => {
+    const { db } = makeDb(undefined);
+    const repo = new AiChatRepo(db);
+
+    await expect(repo.findLatestByPage('u1', 'ws1', 'p1')).resolves.toBeUndefined();
+  });
+});

apps/server/src/database/repos/ai-chat/ai-chat.repo.ts

@@ -80,6 +80,32 @@ export class AiChatRepo {
     });
   }
 
+  /**
+   * The "bound chat" for a document: the requesting user's most recently
+   * created, non-deleted chat whose origin page is `pageId`. Auto-opened when
+   * the AI chat window is opened on that page. Newest-by-createdAt wins, so a
+   * chat created later on the same page supersedes earlier ones — exactly how
+   * "new chat -> becomes the bound one" falls out for free. Scoped to the user +
+   * workspace, so a foreign pageId can only ever match the caller's own chats.
+   */
+  async findLatestByPage(
+    creatorId: string,
+    workspaceId: string,
+    pageId: string,
+  ): Promise<AiChat | undefined> {
+    return this.db
+      .selectFrom('aiChats')
+      .selectAll('aiChats')
+      .where('creatorId', '=', creatorId)
+      .where('workspaceId', '=', workspaceId)
+      .where('pageId', '=', pageId)
+      .where('deletedAt', 'is', null)
+      .orderBy('createdAt', 'desc')
+      .orderBy('id', 'desc') // stable tiebreaker, mirrors findByCreator's cursor
+      .limit(1)
+      .executeTakeFirst();
+  }
+
   async insert(
     insertable: InsertableAiChat,
     trx?: KyselyTransaction,