feat(ai-chat): getCurrentPage отдаёт текущее выделение пользователя (#388) #390
@@ -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 <ChatThread>, 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
|
||||
|
||||
@@ -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(
|
||||
<QueryClientProvider client={queryClient}>
|
||||
<MantineProvider>
|
||||
<ChatThread
|
||||
chatId="c1"
|
||||
initialRows={[]}
|
||||
openPage={props.openPage as never}
|
||||
getEditorSelection={props.getEditorSelection as never}
|
||||
onTurnFinished={vi.fn()}
|
||||
onResumeFallback={vi.fn()}
|
||||
onServerStop={vi.fn()}
|
||||
/>
|
||||
</MantineProvider>
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
}
|
||||
|
||||
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);
|
||||
|
||||
|
||||
@@ -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<OpenPageContext | null>(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,
|
||||
|
||||
@@ -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: "<p>fresh</p>",
|
||||
});
|
||||
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();
|
||||
});
|
||||
});
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
|
||||
@@ -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),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -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 },
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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<Record<string, Tool>> {
|
||||
// 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),
|
||||
}),
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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<string, unknown>;
|
||||
|
||||
// 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,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user