diff --git a/apps/client/src/features/ai-chat/components/ai-chat-window.tsx b/apps/client/src/features/ai-chat/components/ai-chat-window.tsx index 5a918cb9..dc08e756 100644 --- a/apps/client/src/features/ai-chat/components/ai-chat-window.tsx +++ b/apps/client/src/features/ai-chat/components/ai-chat-window.tsx @@ -37,6 +37,14 @@ import { mobileSidebarAtom, } from "@/components/layouts/global/hooks/atoms/sidebar-atom.ts"; import { usePageQuery } from "@/features/page/queries/page-query.ts"; +import { + pageEditorAtom, + readOnlyEditorAtom, +} from "@/features/editor/atoms/editor-atoms.ts"; +import { + getEditorSelectionContext, + type EditorSelectionContext, +} from "@/features/editor/utils/get-editor-selection.ts"; import { extractPageSlugId } from "@/lib"; import { AI_CHATS_RQ_KEY, @@ -314,6 +322,27 @@ export default function AiChatWindow() { ? { id: openPageData.id, title: openPageData.title } : null; + // Live editor handles for the selection snapshot (#388). Both are published by + // the page editor; the read-only editor is used in read mode. Reading the + // selection off `editor.state` stays valid after the editor blurs (ProseMirror + // keeps state.selection), mirroring the comment button (comment-dialog.tsx). + const pageEditor = useAtomValue(pageEditorAtom); + const readOnlyEditor = useAtomValue(readOnlyEditorAtom); + + // Snapshot the user's current editor selection at send time. Edit-mode editor + // wins; the read-only editor is the fallback (read mode). Null when neither + // holds a non-empty selection. Passed to , which reads it live + // from a ref inside prepareSendMessagesRequest — so each turn ships a fresh + // snapshot and multi-turn works without recreating the transport. + const getEditorSelection = useCallback((): EditorSelectionContext | null => { + for (const editor of [pageEditor, readOnlyEditor]) { + if (!editor || editor.isDestroyed) continue; + const sel = getEditorSelectionContext(editor.state); + if (sel) return sel; + } + return null; + }, [pageEditor, readOnlyEditor]); + // The AI-chat thread-identity lifecycle (mount key, both new-chat id adoption // paths, the history-loaded latch, the render-phase reconciler) lives in this // hook. See adopt-chat-id.ts for the canonical #137 two-tab race explanation. @@ -936,6 +965,9 @@ export default function AiChatWindow() { chatId={activeChatId} initialRows={activeChatId ? messageRows : []} openPage={openPage} + // #388: live snapshotter for the user's editor selection, read at + // send time and nested inside openPage on the wire. + getEditorSelection={getEditorSelection} // Honoured only for a new chat; null = universal assistant. roleId={activeChatId === null ? selectedRoleId : null} // Role cards are the new-chat empty-state; offered only when this diff --git a/apps/client/src/features/ai-chat/components/chat-thread.test.tsx b/apps/client/src/features/ai-chat/components/chat-thread.test.tsx index 7d1027c1..20b0a3ec 100644 --- a/apps/client/src/features/ai-chat/components/chat-thread.test.tsx +++ b/apps/client/src/features/ai-chat/components/chat-thread.test.tsx @@ -200,6 +200,74 @@ describe("ChatThread — send now (#198)", () => { }); }); +// #388: the editor selection is snapshotted at send time and nested inside +// openPage on the wire. The getter is read live from a ref, so each send ships a +// fresh snapshot. +describe("ChatThread — editor selection wiring (#388)", () => { + beforeEach(resetState); + afterEach(cleanup); + + function renderWithSelection(props: { + openPage?: { id: string; title: string } | null; + getEditorSelection?: () => unknown; + }) { + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + render( + + + + + , + ); + } + + it("nests the snapshot from the getter into openPage.selection at send time", () => { + const selection = { text: "fix this", blockIds: ["b1"], before: "a " }; + renderWithSelection({ + openPage: { id: "p1", title: "Doc" }, + getEditorSelection: () => selection, + }); + const prep = h.state.transport!.prepareSendMessagesRequest!; + const openPage = prep({ messages: [], body: {} }).body.openPage as Record< + string, + unknown + >; + expect(openPage).toEqual({ id: "p1", title: "Doc", selection }); + }); + + it("sends selection: null when the getter returns null", () => { + renderWithSelection({ + openPage: { id: "p1", title: "Doc" }, + getEditorSelection: () => null, + }); + const prep = h.state.transport!.prepareSendMessagesRequest!; + const openPage = prep({ messages: [], body: {} }).body.openPage as Record< + string, + unknown + >; + expect(openPage).toEqual({ id: "p1", title: "Doc", selection: null }); + }); + + it("does not send selection at all on a non-page route (openPage null)", () => { + const getter = vi.fn(() => ({ text: "sel" })); + renderWithSelection({ openPage: null, getEditorSelection: getter }); + const prep = h.state.transport!.prepareSendMessagesRequest!; + expect(prep({ messages: [], body: {} }).body.openPage).toBeNull(); + // The getter must not even be consulted when there is no page. + expect(getter).not.toHaveBeenCalled(); + }); +}); + describe("ChatThread — turn-end decision (onFinish)", () => { beforeEach(resetState); diff --git a/apps/client/src/features/ai-chat/components/chat-thread.tsx b/apps/client/src/features/ai-chat/components/chat-thread.tsx index acf4326e..5ef54585 100644 --- a/apps/client/src/features/ai-chat/components/chat-thread.tsx +++ b/apps/client/src/features/ai-chat/components/chat-thread.tsx @@ -33,6 +33,7 @@ import { mergeById, } from "@/features/ai-chat/utils/resume-helpers.ts"; import { AI_CHAT_MESSAGES_RQ_KEY } from "@/features/ai-chat/queries/ai-chat-query.ts"; +import type { EditorSelectionContext } from "@/features/editor/utils/get-editor-selection.ts"; import { dequeue, enqueueMessage, @@ -69,6 +70,10 @@ interface ChatThreadProps { /** The page currently open in the workspace, or null on a non-page route. * Sent with each turn so the agent knows what "this page" refers to. */ openPage?: OpenPageContext | null; + /** #388: snapshot the user's current editor selection at SEND time. Invoked + * inside prepareSendMessagesRequest and nested into openPage on the wire, so a + * fresh snapshot ships each turn. Null/absent => nothing selected. */ + getEditorSelection?: () => EditorSelectionContext | null; /** The agent role selected for a NEW chat (null = universal assistant). Sent * in the request body so the server persists it on chat creation; ignored by * the server for existing chats (the role is read from the chat row). */ @@ -151,6 +156,7 @@ export default function ChatThread({ threadKey, initialRows, openPage, + getEditorSelection, roleId, roles, onRolePicked, @@ -226,6 +232,14 @@ export default function ChatThread({ const openPageRef = useRef(openPage ?? null); openPageRef.current = openPage ?? null; + // Keep the selection snapshotter in a ref, same rationale as openPageRef: the + // transport useMemo([]) closes it over, so prop-identity churn must not matter. + // Called at send time inside prepareSendMessagesRequest (#388). + const getEditorSelectionRef = useRef< + (() => EditorSelectionContext | null) | undefined + >(getEditorSelection); + getEditorSelectionRef.current = getEditorSelection; + // Keep the selected role id in a ref, same rationale as openPageRef. Only the // FIRST request of a brand-new chat uses it (the server persists it then and // ignores it for existing chats), but sending it on every send is harmless. @@ -388,7 +402,16 @@ export default function ChatThread({ body: { ...body, chatId: chatIdRef.current, - openPage: openPageRef.current, + // Attach the live editor selection to the open-page context at send + // time — "this"/"here" in the user's message means THIS selection. + // Nested inside openPage so it dies with the page when the server + // rejects the page id (#388). Null when nothing is selected. + openPage: openPageRef.current + ? { + ...openPageRef.current, + selection: getEditorSelectionRef.current?.() ?? null, + } + : null, // Honoured by the server only when creating a new chat; null => // universal assistant. roleId: roleIdRef.current, diff --git a/apps/client/src/features/editor/utils/get-editor-selection.test.ts b/apps/client/src/features/editor/utils/get-editor-selection.test.ts new file mode 100644 index 00000000..d8d6eb62 --- /dev/null +++ b/apps/client/src/features/editor/utils/get-editor-selection.test.ts @@ -0,0 +1,113 @@ +import { describe, it, expect } from "vitest"; +import { Editor } from "@tiptap/core"; +import { Document } from "@tiptap/extension-document"; +import { Paragraph } from "@tiptap/extension-paragraph"; +import { Text } from "@tiptap/extension-text"; +import { EditorState, TextSelection } from "@tiptap/pm/state"; +import type { Node as PMNode } from "@tiptap/pm/model"; +import { UniqueID } from "@docmost/editor-ext"; +import { getEditorSelectionContext } from "./get-editor-selection"; + +/** + * Unit tests for getEditorSelectionContext (#388). Built on a headless + * ProseMirror schema (Document + Paragraph + Text + the block-id UniqueID + * extension), mirroring the editor-ext test style. We assemble docs with + * explicit block ids so the covered-blockIds assertions are deterministic. + */ + +// A schema that carries the `id` block attribute (UniqueID) on paragraphs, just +// like the real editor. +const { schema } = new Editor({ + extensions: [ + Document, + Paragraph, + Text, + UniqueID.configure({ types: ["paragraph"] }), + ], + content: "", +}); + +function docOf(blocks: { id: string; text: string }[]): PMNode { + return schema.node( + "doc", + null, + blocks.map((b) => + schema.node("paragraph", { id: b.id }, b.text ? schema.text(b.text) : []), + ), + ); +} + +function stateWith(doc: PMNode, from: number, to: number): EditorState { + const base = EditorState.create({ schema, doc }); + return base.apply(base.tr.setSelection(TextSelection.create(doc, from, to))); +} + +// Select every text position of the doc (pos 1 .. content.size - 1). +function selectAll(doc: PMNode): EditorState { + return stateWith(doc, 1, doc.content.size - 1); +} + +describe("getEditorSelectionContext", () => { + it("returns null for an empty (collapsed) selection", () => { + const doc = docOf([{ id: "b1", text: "Hello world" }]); + const state = stateWith(doc, 3, 3); // caret, from === to + expect(getEditorSelectionContext(state)).toBeNull(); + }); + + it("returns null for the default caret-at-start of a fresh editor", () => { + const editor = new Editor({ + extensions: [ + Document, + Paragraph, + Text, + UniqueID.configure({ types: ["paragraph"] }), + ], + content: "

fresh

", + }); + expect(getEditorSelectionContext(editor.state)).toBeNull(); + editor.destroy(); + }); + + it("reads a single-paragraph selection with no block-separator artifacts", () => { + const doc = docOf([{ id: "b1", text: "Hello world" }]); + const sel = getEditorSelectionContext(selectAll(doc))!; + expect(sel.text).toBe("Hello world"); + expect(sel.blockIds).toEqual(["b1"]); + expect(sel.truncated).toBeUndefined(); + }); + + it("joins multiple blocks with a newline and collects all covered blockIds", () => { + const doc = docOf([ + { id: "b1", text: "First" }, + { id: "b2", text: "Second" }, + ]); + const sel = getEditorSelectionContext(selectAll(doc))!; + expect(sel.text).toBe("First\nSecond"); + expect(sel.blockIds).toEqual(["b1", "b2"]); + }); + + it("caps the text at 2000 chars and flags truncated", () => { + const doc = docOf([{ id: "b1", text: "x".repeat(2500) }]); + const sel = getEditorSelectionContext(selectAll(doc))!; + expect(sel.text).toHaveLength(2000); + expect(sel.truncated).toBe(true); + }); + + it("computes before/after context and clamps it to the doc bounds", () => { + // One paragraph "0123456789abcdefghij"; select the middle "56789". + const doc = docOf([{ id: "b1", text: "0123456789abcdefghij" }]); + // text char i lives at pos (1 + i); select chars index 5..9 -> pos 6..11. + const sel = getEditorSelectionContext(stateWith(doc, 6, 11))!; + expect(sel.text).toBe("56789"); + expect(sel.before).toBe("01234"); + expect(sel.after).toBe("abcdefghij"); + }); + + it("omits before/after at the document boundaries (never reads past 0/size)", () => { + const doc = docOf([{ id: "b1", text: "Edge" }]); + const sel = getEditorSelectionContext(selectAll(doc))!; + // Selection spans the whole single block: nothing before or after it. + expect(sel.before).toBeUndefined(); + expect(sel.after).toBeUndefined(); + }); +}); diff --git a/apps/client/src/features/editor/utils/get-editor-selection.ts b/apps/client/src/features/editor/utils/get-editor-selection.ts new file mode 100644 index 00000000..d711a335 --- /dev/null +++ b/apps/client/src/features/editor/utils/get-editor-selection.ts @@ -0,0 +1,71 @@ +import type { EditorState } from "@tiptap/pm/state"; + +export interface EditorSelectionContext { + text: string; + truncated?: boolean; + blockIds?: string[]; + before?: string; + after?: string; +} + +// Client-side caps. The server re-caps every field independently (defence in +// depth — the payload is attacker-controllable), so these only keep the wire +// small for the common case. +const TEXT_CAP = 2000; +const CONTEXT_CHARS = 160; +const MAX_BLOCK_IDS = 20; + +// Pure: takes an EditorState so it is unit-testable with a headless editor. +// Snapshots the user's current selection into the wire shape carried inside +// openPage — plain text + the ids of the blocks it covers + a little surrounding +// context. Returns null when nothing meaningful is selected. +// +// Deliberately does NOT emit the ProseMirror positions (from/to): they rot the +// instant the document changes and the server tools address content by block id +// + text (getNode / editPageText find-replace), never by position. +export function getEditorSelectionContext( + state: EditorState, +): EditorSelectionContext | null { + const { selection, doc } = state; + // An empty selection (incl. the default caret-at-start of a fresh editor) is + // never a "this"/"here" — bail before reading any text. + if (selection.empty) return null; + + const { from, to } = selection; + + let text = doc.textBetween(from, to, "\n"); + let truncated = false; + if (text.length > TEXT_CAP) { + text = text.slice(0, TEXT_CAP); + truncated = true; + } + // A selection spanning only non-text nodes (e.g. an image) trims to empty -> + // treat as no selection. + if (text.trim().length === 0) return null; + + // Ids of every block the selection covers, deduped and capped. These bridge + // the plain-text selection to the server tools (getNode / editPageText). + const blockIds: string[] = []; + doc.nodesBetween(from, to, (node) => { + const id = node.isBlock ? node.attrs?.id : undefined; + if (typeof id === "string" && id.length > 0 && !blockIds.includes(id)) { + blockIds.push(id); + } + }); + + // ~160 chars of plain text on each side, clamped to the document bounds, so + // editPageText can disambiguate a duplicate of the selected text. + const before = doc.textBetween(Math.max(0, from - CONTEXT_CHARS), from, "\n"); + const after = doc.textBetween( + to, + Math.min(doc.content.size, to + CONTEXT_CHARS), + "\n", + ); + + const result: EditorSelectionContext = { text }; + if (truncated) result.truncated = true; + if (blockIds.length > 0) result.blockIds = blockIds.slice(0, MAX_BLOCK_IDS); + if (before.length > 0) result.before = before; + if (after.length > 0) result.after = after; + return result; +} diff --git a/apps/server/src/core/ai-chat/ai-chat.prompt.spec.ts b/apps/server/src/core/ai-chat/ai-chat.prompt.spec.ts index 9abd7de1..4186ff51 100644 --- a/apps/server/src/core/ai-chat/ai-chat.prompt.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.prompt.spec.ts @@ -153,6 +153,41 @@ describe('buildSystemPrompt current-page context', () => { expect(prompt).not.toContain('pageId:'); }); + // #388: editor-selection flag. Only a FIXED one-liner is added — the selection + // TEXT (untrusted page content) must never reach the prompt. + const SELECTION_FLAG = 'currently has text SELECTED on this page'; + + it('adds the selection flag when a selection is present with a page', () => { + const prompt = buildSystemPrompt({ + workspace, + openedPage: { + id: 'pg-123', + title: 'Doc', + selection: { text: 'SECRET-SELECTED-TEXT', blockIds: ['b1'] }, + }, + }); + expect(prompt).toContain(SELECTION_FLAG); + // The selection TEXT itself is NEVER in the prompt. + expect(prompt).not.toContain('SECRET-SELECTED-TEXT'); + expect(prompt).not.toContain('b1'); + }); + + it('omits the selection flag when there is no selection', () => { + const prompt = buildSystemPrompt({ + workspace, + openedPage: { id: 'pg-123', title: 'Doc' }, + }); + expect(prompt).not.toContain(SELECTION_FLAG); + }); + + it('omits the selection flag when selection is null', () => { + const prompt = buildSystemPrompt({ + workspace, + openedPage: { id: 'pg-123', title: 'Doc', selection: null }, + }); + expect(prompt).not.toContain(SELECTION_FLAG); + }); + it('escapes a malicious opened-page title so it cannot inject tags (F1)', () => { const prompt = buildSystemPrompt({ workspace, diff --git a/apps/server/src/core/ai-chat/ai-chat.prompt.ts b/apps/server/src/core/ai-chat/ai-chat.prompt.ts index 337f60ca..3e7cb876 100644 --- a/apps/server/src/core/ai-chat/ai-chat.prompt.ts +++ b/apps/server/src/core/ai-chat/ai-chat.prompt.ts @@ -156,8 +156,13 @@ export interface BuildSystemPromptInput { * has an id, a CONTEXT line is added so the agent can resolve "this page" / * "the current page" to that pageId. The page is NOT fetched here — the agent * uses its CASL-enforced read/write page tools with the id when needed. + * + * `selection` (#388) is present only when the user has a non-empty editor + * selection; the prompt adds ONLY a fixed one-line flag from it — the + * selection TEXT is untrusted page content and stays out of the prompt (it is + * surfaced solely via the getCurrentPage tool result). */ - openedPage?: { id?: string; title?: string } | null; + openedPage?: { id?: string; title?: string; selection?: object | null } | null; /** * Admin-authored, per-EXTERNAL-MCP-server guidance ("how/when to use this * server's tools"), built by `McpClientsService.toolsFor` for servers that @@ -309,6 +314,14 @@ export function buildSystemPrompt({ ? escapeAttr(openedPage.title) : 'Untitled'; context += `\nThe user is currently viewing the page "${title}" (pageId: ${pageId.trim()}). When they refer to "this page", "the current page", or similar, operate on that pageId — use the read/write page tools with it.`; + // Editor-selection flag (#388). A FIXED one-liner only — the selection TEXT + // is untrusted collaborative-page content and must never enter the prompt; it + // is surfaced solely through the getCurrentPage tool result (SAFETY_FRAMEWORK + // treats a tool result as data). Nested under the page block so it is added + // only alongside a resolved page (a selection cannot outlive its page). + if (openedPage?.selection) { + context += `\nThe user currently has text SELECTED on this page — call getCurrentPage to see the selection. When they say "this", "here", "the selected text" or similar, they mean that selection.`; + } } // Interrupt-resume marker (#198). Added to the context section (inside the diff --git a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts index 19e82cee..31c985c4 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts @@ -698,6 +698,7 @@ describe('AiChatService.resolveOpenPageContext (#159 current-page validation)', id: string; title: string; updatedAt: Date; + selection: unknown; } | null>; it('returns null when no page is open (no id)', async () => { @@ -743,8 +744,14 @@ describe('AiChatService.resolveOpenPageContext (#159 current-page validation)', }); // The client claims it is on "Page A" but the id points at page B. const result = await call(svc, { id: 'p-1', title: 'Page A' }); - // updatedAt (#274 page-change fast path) is carried through from the DB row. - expect(result).toEqual({ id: 'p-1', title: 'Real Title B', updatedAt }); + // updatedAt (#274 page-change fast path) is carried through from the DB row; + // selection is null when the client sent none (#388). + expect(result).toEqual({ + id: 'p-1', + title: 'Real Title B', + updatedAt, + selection: null, + }); }); it('coerces a null DB title to an empty string', async () => { @@ -757,8 +764,55 @@ describe('AiChatService.resolveOpenPageContext (#159 current-page validation)', id: 'p-1', title: '', updatedAt, + selection: null, }); }); + + // #388: the selection rides ONLY on a successful page resolve, and is + // sanitized on the way through. + it('attaches the SANITIZED selection on a successful resolve', async () => { + const updatedAt = new Date('2026-07-02T10:00:00Z'); + const svc = makeService({ + page: { id: 'p-1', workspaceId: 'ws-1', title: 'Doc', updatedAt }, + canView: true, + }); + const result = await call(svc, { + id: 'p-1', + selection: { + text: 'fix this', + blockIds: ['b1', 123, 'y'.repeat(65)], // garbage stripped by sanitize + before: 'please ', + }, + }); + expect(result).toEqual({ + id: 'p-1', + title: 'Doc', + updatedAt, + selection: { text: 'fix this', blockIds: ['b1'], before: 'please ' }, + }); + }); + + it('drops a blank/garbage selection to null on a successful resolve', async () => { + const updatedAt = new Date('2026-07-02T10:00:00Z'); + const svc = makeService({ + page: { id: 'p-1', workspaceId: 'ws-1', title: 'Doc', updatedAt }, + canView: true, + }); + expect( + await call(svc, { id: 'p-1', selection: { text: ' ' } }), + ).toEqual({ id: 'p-1', title: 'Doc', updatedAt, selection: null }); + }); + + it('selection does NOT survive a foreign/inaccessible page (dies with the page)', async () => { + // Forbidden page => the WHOLE context is null, so the selection is gone too. + const svc = makeService({ + page: { id: 'p-1', workspaceId: 'ws-1', title: 'Restricted' }, + canView: false, + }); + expect( + await call(svc, { id: 'p-1', selection: { text: 'secret sel' } }), + ).toBeNull(); + }); }); /** diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index d8e027b9..7a1a8af3 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -43,6 +43,10 @@ import { } from './tools/tool-tiers'; import { RunAlreadyActiveError } from './ai-chat-run.service'; import { computePageChange } from './page-change/page-change.util'; +import { + sanitizeSelection, + type SelectionContext, +} from './tools/current-page.util'; import { roleModelOverride } from './roles/role-model-config'; import { startSseHeartbeat, @@ -189,7 +193,12 @@ export interface AiChatStreamBody { // page" refers to; the page itself is never fetched server-side here. The id // is attacker-controllable but harmless: the agent reads/writes via its // CASL-enforced page tools, which 403 on a page the user cannot access. - openPage?: { id?: string; title?: string } | null; + // + // `selection` is the user's editor selection snapshotted client-side at send + // time (#388). It is CLIENT-controlled and UNTRUSTED — a loose `unknown` here + // (the body is parsed off req.body without a DTO) that is type-checked and + // capped by `sanitizeSelection` before it is ever surfaced to the model. + openPage?: { id?: string; title?: string; selection?: unknown } | null; // Set by the client "send now" action (#198): this turn immediately follows a // user interruption of the previous turn. A hint only — the server re-confirms // it against persisted history (`isInterruptResume`) before injecting the @@ -365,10 +374,18 @@ export class AiChatService implements OnModuleInit { * page, or any non-Forbidden access-check fault, returns null. */ private async resolveOpenPageContext( - openPage: { id?: string; title?: string } | null | undefined, + openPage: + | { id?: string; title?: string; selection?: unknown } + | null + | undefined, workspace: Workspace, user: User, - ): Promise<{ id: string; title: string; updatedAt: Date } | null> { + ): Promise<{ + id: string; + title: string; + updatedAt: Date; + selection: SelectionContext | null; + } | null> { const candidatePageId = openPage?.id; if (!candidatePageId) return null; const page = await this.pageRepo.findById(candidatePageId); @@ -390,7 +407,18 @@ export class AiChatService implements OnModuleInit { // updatedAt is the page's last-modified instant, used by the #274 per-turn // page-change detection as a cheap fast path (unchanged instant => skip the // render + diff). The system-prompt / tool consumers ignore the extra field. - return { id: page.id, title: page.title ?? '', updatedAt: page.updatedAt }; + // + // The sanitized editor selection (#388) is attached ONLY here, on a + // successful page resolve: the fail-closed branches above return null for the + // WHOLE context, so a selection can never outlive a foreign/missing/deleted + // page (decision 3). Downstream consumers that don't care (detectPageChange, + // snapshotOpenPage) ignore the extra field, same as updatedAt. + return { + id: page.id, + title: page.title ?? '', + updatedAt: page.updatedAt, + selection: sanitizeSelection(openPage?.selection), + }; } /** diff --git a/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.spec.ts b/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.spec.ts index 7e4b35d6..fd5ef4c2 100644 --- a/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.spec.ts +++ b/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.spec.ts @@ -651,3 +651,69 @@ describe('AiChatToolsService #294 changed execute wirings', () => { expect(calls.tableUpdateCell).toEqual([['p1', '#0', 1, 2, 'x']]); }); }); + +/** + * getCurrentPage selection contract (#388): the tool surfaces the selection that + * was sanitized + nested onto the resolved open-page context (last forUser arg). + * No page => selection is null. The tool never fetches or verifies anything — it + * just projects the resolved context. + */ +describe('AiChatToolsService getCurrentPage selection (#388)', () => { + const tokenServiceStub = { + generateAccessToken: jest.fn().mockResolvedValue('access-token'), + generateCollabToken: jest.fn().mockResolvedValue('collab-token'), + }; + + let service: AiChatToolsService; + + beforeEach(() => { + jest.spyOn(loader, 'loadDocmostMcp').mockResolvedValue( + mockLoaded(function () { + return {} as DocmostClientLike; + } as unknown as loader.DocmostClientCtor), + ); + service = new AiChatToolsService( + tokenServiceStub as never, + {} as never, + {} as never, + {} as never, + {} as never, + { + asSink: () => ({ put: jest.fn(), has: jest.fn(), evict: jest.fn() }), + } as never, + ); + }); + + afterEach(() => jest.restoreAllMocks()); + + const buildTools = (openedPage: unknown) => + service.forUser( + { id: 'user-1', email: 'u@example.com', workspaceId: 'ws-1' } as never, + 'session-1', + 'ws-1', + 'chat-1', + openedPage as never, + ); + + it('returns the nested selection from the resolved context', async () => { + const selection = { text: 'fix this', blockIds: ['b1'], before: 'a ' }; + const tools = await buildTools({ id: 'p1', title: 'Doc', selection }); + expect(await tools.getCurrentPage.execute({} as never, {} as never)).toEqual( + { page: { id: 'p1', title: 'Doc' }, selection }, + ); + }); + + it('returns selection: null when the context has no selection', async () => { + const tools = await buildTools({ id: 'p1', title: 'Doc' }); + expect(await tools.getCurrentPage.execute({} as never, {} as never)).toEqual( + { page: { id: 'p1', title: 'Doc' }, selection: null }, + ); + }); + + it('returns { page: null, selection: null } when no page is open', async () => { + const tools = await buildTools(null); + expect(await tools.getCurrentPage.execute({} as never, {} as never)).toEqual( + { page: null, selection: null }, + ); + }); +}); diff --git a/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts b/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts index d543390e..29b427ea 100644 --- a/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts +++ b/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.ts @@ -13,7 +13,10 @@ import { type DocmostClientLike, type SharedToolSpec, } from './docmost-client.loader'; -import { resolveCurrentPageResult } from './current-page.util'; +import { + resolveCurrentPageResult, + type SelectionContext, +} from './current-page.util'; import { parseNodeArg } from './parse-node-arg'; import { modelFriendlyInput } from './model-friendly-input'; import { SandboxStore } from '../../../integrations/sandbox/sandbox.store'; @@ -153,8 +156,13 @@ export class AiChatToolsService { // The page the user currently has open (from the request context), exposed // to the model via getCurrentPage. Optional and last so existing callers // keep compiling. Kept proxy-robust: the model can CALL for the current - // page instead of relying on it surviving in the system prompt text. - openedPage?: { id?: string; title?: string } | null, + // page instead of relying on it surviving in the system prompt text. The + // `selection` (#388) is already sanitized + nested by resolveOpenPageContext. + openedPage?: { + id?: string; + title?: string; + selection?: SelectionContext | null; + } | null, ): Promise> { // Build the per-user loopback client (carrying the access + collab // provenance tokens) and load the shared tool-spec registry. Client @@ -309,9 +317,15 @@ export class AiChatToolsService { getCurrentPage: tool({ description: 'Return the page the user is currently viewing — i.e. what "this page", ' + - '"the current page", or "here" refers to. Returns the page id and title, ' + - 'or null if the user is not currently on a page. Call this first whenever ' + - 'the user refers to the current page without giving an explicit id.', + '"the current page", or "here" refers to — plus the text the user ' + + 'currently has SELECTED on that page (what "this", "here", "the selected ' + + 'fragment" refers to), or selection: null when nothing is selected. The ' + + 'selection is a client-side snapshot taken when the user sent the message ' + + 'and includes the ids of the blocks it covers plus surrounding context; ' + + 'it is NOT verified server-side — locate it in the page (searchInPage / ' + + 'getNode) before editing. Returns page: null if the user is not currently ' + + 'on a page. Call this first whenever the user refers to the current page ' + + 'or a selected fragment without giving an explicit id.', inputSchema: modelFriendlyInput({}), execute: async () => resolveCurrentPageResult(openedPage), }), diff --git a/apps/server/src/core/ai-chat/tools/current-page.util.spec.ts b/apps/server/src/core/ai-chat/tools/current-page.util.spec.ts index d0649773..19af9373 100644 --- a/apps/server/src/core/ai-chat/tools/current-page.util.spec.ts +++ b/apps/server/src/core/ai-chat/tools/current-page.util.spec.ts @@ -1,43 +1,180 @@ -import { resolveCurrentPageResult } from './current-page.util'; +import { + resolveCurrentPageResult, + sanitizeSelection, +} from './current-page.util'; /** * Unit tests for resolveCurrentPageResult (pure function). Mirrors the - * getCurrentPage tool's contract: { page: null } when no page is open (no id), - * otherwise { page: { id, title } } with title defaulting to ''. + * getCurrentPage tool's contract: { page: null, selection: null } when no page + * is open (no id), otherwise { page: { id, title }, selection } with title + * defaulting to '' and the selection passed through from the resolved context. */ describe('resolveCurrentPageResult', () => { - it('returns { page: null } when openedPage is undefined', () => { - expect(resolveCurrentPageResult(undefined)).toEqual({ page: null }); + it('returns { page: null, selection: null } when openedPage is undefined', () => { + expect(resolveCurrentPageResult(undefined)).toEqual({ + page: null, + selection: null, + }); }); - it('returns { page: null } when openedPage is null', () => { - expect(resolveCurrentPageResult(null)).toEqual({ page: null }); + it('returns { page: null, selection: null } when openedPage is null', () => { + expect(resolveCurrentPageResult(null)).toEqual({ + page: null, + selection: null, + }); }); - it('returns { page: null } when openedPage has no id', () => { - expect(resolveCurrentPageResult({})).toEqual({ page: null }); - expect(resolveCurrentPageResult({ title: 'x' })).toEqual({ page: null }); + it('returns { page: null, selection: null } when openedPage has no id', () => { + expect(resolveCurrentPageResult({})).toEqual({ + page: null, + selection: null, + }); + expect(resolveCurrentPageResult({ title: 'x' })).toEqual({ + page: null, + selection: null, + }); }); - it('returns { page: null } when id is an empty string', () => { - expect(resolveCurrentPageResult({ id: '' })).toEqual({ page: null }); + it('returns { page: null, selection: null } when id is an empty string', () => { + expect(resolveCurrentPageResult({ id: '' })).toEqual({ + page: null, + selection: null, + }); }); - it('returns the page id and title when both are present', () => { + it('drops the selection when there is no page (selection dies with the page)', () => { + // Even if a selection somehow rode along without a page id, a null page + // always yields a null selection. + expect( + resolveCurrentPageResult({ selection: { text: 'orphan' } }), + ).toEqual({ page: null, selection: null }); + }); + + it('returns the page id and title with a null selection by default', () => { expect(resolveCurrentPageResult({ id: 'p1', title: 'Hello' })).toEqual({ page: { id: 'p1', title: 'Hello' }, + selection: null, + }); + }); + + it('passes the nested selection through verbatim', () => { + const selection = { + text: 'fix this', + blockIds: ['b1'], + before: 'please ', + after: ' now', + }; + expect( + resolveCurrentPageResult({ id: 'p1', title: 'Hello', selection }), + ).toEqual({ + page: { id: 'p1', title: 'Hello' }, + selection, }); }); it('defaults title to "" when it is missing', () => { expect(resolveCurrentPageResult({ id: 'p1' })).toEqual({ page: { id: 'p1', title: '' }, + selection: null, }); }); it('keeps an explicit empty-string title as ""', () => { expect(resolveCurrentPageResult({ id: 'p1', title: '' })).toEqual({ page: { id: 'p1', title: '' }, + selection: null, }); }); }); + +/** + * Unit tests for sanitizeSelection (#388). The selection is an attacker- + * controllable client snapshot: every field is type-checked and capped, and + * anything that is not a real selection collapses to null. It is NEVER verified + * against the page content (decision 5 — a hint, not ground truth). + */ +describe('sanitizeSelection', () => { + it('accepts a well-formed payload unchanged', () => { + const raw = { + text: 'the selected fragment', + truncated: true, + blockIds: ['b1', 'b2'], + before: 'context before ', + after: ' context after', + }; + expect(sanitizeSelection(raw)).toEqual(raw); + }); + + it('returns null for non-objects', () => { + expect(sanitizeSelection(null)).toBeNull(); + expect(sanitizeSelection(undefined)).toBeNull(); + expect(sanitizeSelection('text')).toBeNull(); + expect(sanitizeSelection(42)).toBeNull(); + expect(sanitizeSelection([])).toBeNull(); + }); + + it('returns null when text is missing, non-string or blank-after-trim', () => { + expect(sanitizeSelection({})).toBeNull(); + expect(sanitizeSelection({ text: 123 })).toBeNull(); + expect(sanitizeSelection({ text: '' })).toBeNull(); + expect(sanitizeSelection({ text: ' \n ' })).toBeNull(); + }); + + it('keeps only text when the other fields are garbage', () => { + expect( + sanitizeSelection({ + text: 'hello', + truncated: 'yes', + blockIds: 'nope', + before: 5, + after: {}, + }), + ).toEqual({ text: 'hello' }); + }); + + it('caps text at 4000 and forces truncated', () => { + const raw = { text: 'a'.repeat(5000) }; + const out = sanitizeSelection(raw)!; + expect(out.text).toHaveLength(4000); + expect(out.truncated).toBe(true); + }); + + it('does not set truncated for text under the cap', () => { + expect(sanitizeSelection({ text: 'short' })).toEqual({ text: 'short' }); + }); + + it('slices blockIds to 20 and drops non-string / oversized ids', () => { + const ids = Array.from({ length: 30 }, (_, i) => `b${i}`); + const out = sanitizeSelection({ + text: 'x', + blockIds: [...ids, 123, '', 'y'.repeat(65)], + })!; + // The 30 valid ids cap to the first 20; the number, empty string and the + // 65-char id are dropped before the slice. + expect(out.blockIds).toHaveLength(20); + expect(out.blockIds).toEqual(ids.slice(0, 20)); + }); + + it('keeps a 64-char id but drops a 65-char one (boundary)', () => { + const ok = 'z'.repeat(64); + const tooLong = 'z'.repeat(65); + expect( + sanitizeSelection({ text: 'x', blockIds: [ok, tooLong] })!.blockIds, + ).toEqual([ok]); + }); + + it('omits blockIds entirely when none survive', () => { + const out = sanitizeSelection({ text: 'x', blockIds: [123, ''] })!; + expect(out.blockIds).toBeUndefined(); + }); + + it('caps before/after at 200 chars and drops empty ones', () => { + const out = sanitizeSelection({ + text: 'x', + before: 'b'.repeat(300), + after: '', + })!; + expect(out.before).toHaveLength(200); + expect(out.after).toBeUndefined(); + }); +}); diff --git a/apps/server/src/core/ai-chat/tools/current-page.util.ts b/apps/server/src/core/ai-chat/tools/current-page.util.ts index 0ced2492..d804daab 100644 --- a/apps/server/src/core/ai-chat/tools/current-page.util.ts +++ b/apps/server/src/core/ai-chat/tools/current-page.util.ts @@ -1,21 +1,91 @@ +export interface SelectionContext { + text: string; + truncated?: boolean; + blockIds?: string[]; + before?: string; + after?: string; +} + +// Server-side caps for the client-reported selection. Intentionally >= the +// client caps: the client pre-trims for a small wire, but this layer re-checks +// everything because the payload is attacker-controllable. +const TEXT_CAP = 4000; +const CONTEXT_CAP = 200; +const MAX_BLOCK_IDS = 20; +const BLOCK_ID_CAP = 64; + +// Sanitize the client-reported selection: type-check every field, cap sizes +// (text 4000, before/after 200, blockIds 20 x 64 chars), drop garbage to null. +// The selection is a CLIENT-side snapshot — never verified against the page +// content (#159 lesson: treat as a hint, not ground truth). The agent is told +// (getCurrentPage's description) to localize it before editing. +export function sanitizeSelection(raw: unknown): SelectionContext | null { + if (!raw || typeof raw !== 'object') return null; + const r = raw as Record; + + // text is the only required field; anything else is a non-selection. + if (typeof r.text !== 'string' || r.text.trim().length === 0) return null; + let text = r.text; + let truncated = r.truncated === true; + if (text.length > TEXT_CAP) { + text = text.slice(0, TEXT_CAP); + truncated = true; + } + + const result: SelectionContext = { text }; + if (truncated) result.truncated = true; + + if (Array.isArray(r.blockIds)) { + // Keep only well-formed, in-range ids (an oversize id is DROPPED, not + // truncated — a mangled id is worse than a missing one), then cap the count. + const ids = r.blockIds + .filter( + (x): x is string => + typeof x === 'string' && x.length > 0 && x.length <= BLOCK_ID_CAP, + ) + .slice(0, MAX_BLOCK_IDS); + if (ids.length > 0) result.blockIds = ids; + } + + if (typeof r.before === 'string' && r.before.length > 0) { + result.before = r.before.slice(0, CONTEXT_CAP); + } + if (typeof r.after === 'string' && r.after.length > 0) { + result.after = r.after.slice(0, CONTEXT_CAP); + } + + return result; +} + export interface CurrentPageInput { id?: string; title?: string; + // The already-sanitized selection nested onto the resolved open-page context + // by resolveOpenPageContext (never the raw client value). Passed through to + // the tool result verbatim; null when nothing is selected. + selection?: SelectionContext | null; } export interface CurrentPageResult { page: { id: string; title: string } | null; + selection: SelectionContext | null; // null when nothing is selected or no page } // Resolve the "current page" tool result from the client-supplied open-page -// context. Returns { page: null } when no page is open (no id), otherwise the -// page id + title (title defaults to '' when absent). Mirrors the getCurrentPage -// tool's contract so it can be unit-tested without the ESM Docmost client. +// context. Returns { page: null, selection: null } when no page is open (no id), +// otherwise the page id + title (title defaults to '' when absent) plus the +// selection already sanitized+nested by resolveOpenPageContext. A null page +// always yields a null selection (the selection dies with the page). Mirrors the +// getCurrentPage tool's contract so it can be unit-tested without the ESM +// Docmost client. export function resolveCurrentPageResult( openedPage?: CurrentPageInput | null, ): CurrentPageResult { if (!openedPage?.id) { - return { page: null }; + return { page: null, selection: null }; } - return { page: { id: openedPage.id, title: openedPage.title ?? '' } }; + return { + page: { id: openedPage.id, title: openedPage.title ?? '' }, + selection: openedPage.selection ?? null, + }; } diff --git a/apps/server/src/core/ai-chat/tools/tool-tiers.ts b/apps/server/src/core/ai-chat/tools/tool-tiers.ts index 7d26498c..f4819fb2 100644 --- a/apps/server/src/core/ai-chat/tools/tool-tiers.ts +++ b/apps/server/src/core/ai-chat/tools/tool-tiers.ts @@ -98,7 +98,8 @@ export const INLINE_TOOL_TIERS: Record< }, getCurrentPage: { tier: 'core', - catalogLine: 'getCurrentPage — the page the user is currently viewing.', + catalogLine: + 'getCurrentPage — the page the user is currently viewing and their current text selection on it.', }, // NOTE: getPage and listPages moved to @docmost/mcp's SHARED_TOOL_SPECS // (#294); they carry their own tier ('core') + catalogLine there.