feat(ai-chat): getCurrentPage отдаёт текущее выделение пользователя (#388) #390

Merged
vvzvlad merged 1 commits from feat/388-getcurrentpage-selection into develop 2026-07-06 19:08:29 +03:00
14 changed files with 758 additions and 33 deletions
@@ -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,
+14 -1
View File
@@ -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.