Merge branch 'develop' into feat/ai-chat-review-followups

Integrate the already-merged step-limit work from develop. Only conflict was
ai-chat.service.spec.ts: both sides appended a describe block and edited the
import line. Resolved as a union — keep compactToolOutput + the assistantParts/
serializeSteps/rowToUiMessage suites (this branch) AND the prepareAgentStep
suite (develop), importing all symbols from ai-chat.service.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
vvzvlad
2026-06-20 18:09:17 +03:00
26 changed files with 895 additions and 1183 deletions

View File

@@ -3,6 +3,9 @@ import {
assistantParts,
serializeSteps,
rowToUiMessage,
prepareAgentStep,
MAX_AGENT_STEPS,
FINAL_STEP_INSTRUCTION,
} from './ai-chat.service';
import type { AiChatMessage } from '@docmost/db/types/entity.types';
@@ -190,3 +193,39 @@ describe('rowToUiMessage', () => {
expect(ui.parts).toEqual([{ type: 'text', text: 'hi there' }]);
});
});
/**
* Unit tests for prepareAgentStep: the pure helper that decides per-step
* overrides for the agent loop. Early steps return undefined (default
* behavior); the final allowed step (stepNumber === MAX_AGENT_STEPS - 1) forces
* a text-only synthesis answer (toolChoice 'none') with the FINAL_STEP_INSTRUCTION
* appended onto — not replacing — the original system prompt.
*/
describe('prepareAgentStep', () => {
it('returns undefined for the first step', () => {
expect(prepareAgentStep(0, 'SYS')).toBeUndefined();
});
it('returns undefined for a non-final step (just before the last)', () => {
expect(prepareAgentStep(MAX_AGENT_STEPS - 2, 'SYS')).toBeUndefined();
});
it('forces a text-only synthesis on the final allowed step', () => {
const result = prepareAgentStep(MAX_AGENT_STEPS - 1, 'SYS');
expect(result).toBeDefined();
expect(result?.toolChoice).toBe('none');
// The original persona is preserved (prefix), not replaced.
expect(result?.system.startsWith('SYS')).toBe(true);
// The synthesis instruction is appended.
expect(result?.system).toContain(FINAL_STEP_INSTRUCTION);
});
it('pins the off-by-one boundary (MAX-2 is not final, MAX-1 is)', () => {
// Boundary expressed via the constant, not a hardcoded 18/19, so the test
// tracks MAX_AGENT_STEPS if the cap ever changes.
expect(prepareAgentStep(MAX_AGENT_STEPS - 2, 'SYS')).toBeUndefined();
const atBoundary = prepareAgentStep(MAX_AGENT_STEPS - 1, 'SYS');
expect(atBoundary).toBeDefined();
expect(atBoundary?.toolChoice).toBe('none');
});
});

View File

@@ -18,6 +18,42 @@ import { AiChatToolsService } from './tools/ai-chat-tools.service';
import { McpClientsService } from './external-mcp/mcp-clients.service';
import { buildSystemPrompt } from './ai-chat.prompt';
// Max agent steps per turn. One step = one model generation; a step that calls
// tools is followed by another step carrying the tool results. Raised from 8 so
// multi-search research questions are not cut off mid-investigation.
const MAX_AGENT_STEPS = 20;
// System-prompt addendum injected ONLY on the final step (see prepareAgentStep).
// It forbids further tool calls and tells the model to synthesize the best
// answer it can from what it already gathered, so a tool-heavy turn never ends
// empty.
const FINAL_STEP_INSTRUCTION =
'You have reached the maximum number of tool-use steps for this turn. ' +
'Do NOT call any more tools. Using only the information already gathered, ' +
"write the most complete, useful final answer you can now, in the user's " +
'language. If the information is incomplete, say so explicitly: summarize ' +
'what you found, what is still missing, and give your best partial conclusion.';
// Pure, unit-testable: decide per-step overrides. Returns undefined for normal
// steps; on the final allowed step forces a text-only synthesis answer.
// `system` is the in-scope system prompt; we CONCATENATE so the original
// persona/context is preserved — a bare `system` override would REPLACE the
// whole system prompt for the step.
//
// NOTE: at AI SDK v7 the per-step `system` field is renamed to `instructions`.
// On v6 (`^6.0.134`) `system` is the correct field — adjust when bumping.
export function prepareAgentStep(
stepNumber: number,
system: string,
): { toolChoice: 'none'; system: string } | undefined {
if (stepNumber >= MAX_AGENT_STEPS - 1) {
return { toolChoice: 'none', system: `${system}\n\n${FINAL_STEP_INSTRUCTION}` };
}
return undefined;
}
export { MAX_AGENT_STEPS, FINAL_STEP_INSTRUCTION };
/**
* Payload accepted from the client `useChat` POST body. We do NOT bind a strict
* DTO (the global ValidationPipe whitelist would strip the useChat-specific
@@ -245,7 +281,13 @@ export class AiChatService {
// cap would truncate complex tool calls mid-argument. Let the model use its
// natural per-step budget. (Cost/credit limits are an account concern, not
// something to enforce by silently breaking the agent.)
stopWhen: stepCountIs(8),
stopWhen: stepCountIs(MAX_AGENT_STEPS),
// Forced finalization: reserve the LAST allowed step for a text-only
// answer. Without this, a turn that spends all its steps on tool calls
// ends with no assistant text (an empty turn). prepareAgentStep forbids
// further tool calls and appends a synthesis instruction on that step,
// concatenated onto the original `system` so the persona is preserved.
prepareStep: ({ stepNumber }) => prepareAgentStep(stepNumber, system),
abortSignal: signal,
onFinish: async ({ text, finishReason, totalUsage, usage, steps }) => {
await persistAssistant({

View File

@@ -578,6 +578,49 @@ export class PageController {
);
}
@HttpCode(HttpStatus.OK)
@Post('/tree')
async getPagesTree(
@Body() dto: SidebarPageDto,
@AuthUser() user: User,
) {
if (!dto.spaceId && !dto.pageId) {
throw new BadRequestException(
'Either spaceId or pageId must be provided',
);
}
let spaceId = dto.spaceId;
if (dto.pageId) {
const page = await this.pageRepo.findById(dto.pageId);
if (!page) {
throw new ForbiddenException();
}
spaceId = page.spaceId;
}
const ability = await this.spaceAbility.createForUser(user, spaceId);
if (ability.cannot(SpaceCaslAction.Read, SpaceCaslSubject.Page)) {
throw new ForbiddenException();
}
const spaceCanEdit = ability.can(
SpaceCaslAction.Edit,
SpaceCaslSubject.Page,
);
const items = await this.pageService.getSidebarPagesTree(
spaceId,
user.id,
spaceCanEdit,
dto.pageId,
);
return { items };
}
@HttpCode(HttpStatus.OK)
@Post('move-to-space')
async movePageToSpace(

View File

@@ -1137,7 +1137,7 @@ export class PageService {
T extends { id: string; parentPageId: string | null },
>(
pages: T[],
rootPageId: string,
rootPageId: string | null,
userId: string,
spaceId?: string,
): Promise<T[]> {
@@ -1153,6 +1153,15 @@ export class PageService {
);
const accessibleSet = new Set(accessibleIds);
// When no explicit root is given (whole-space tree), every page whose
// parent is outside the returned set acts as a root (space root pages have
// parentPageId === null). This mirrors the single-root case below.
const pageIdSet = new Set(pageIds);
const isRoot = (page: T): boolean => {
if (rootPageId !== null) return page.id === rootPageId;
return !page.parentPageId || !pageIdSet.has(page.parentPageId);
};
// Prune: include a page only if it's accessible AND its parent chain to root is included
const includedIds = new Set<string>();
@@ -1166,7 +1175,7 @@ export class PageService {
if (!accessibleSet.has(page.id)) continue;
// Root page: include if accessible
if (page.id === rootPageId) {
if (isRoot(page)) {
includedIds.add(page.id);
changed = true;
continue;
@@ -1182,4 +1191,123 @@ export class PageService {
return pages.filter((p) => includedIds.has(p.id));
}
/**
* Whole subtree (pageId) or whole space tree (spaceId only) in a single
* query, permission-filtered, returned as a flat list matching the sidebar
* item shape (id, slugId, title, icon, position, parentPageId, spaceId,
* hasChildren, canEdit) ordered by position. content is never fetched.
*
* Reproduces the exact two-branch permission logic of getSidebarPages():
* - open space (no restrictions): every returned page is visible, canEdit =
* spaceCanEdit, hasChildren derived from the returned set.
* - restricted space: full descendant set is loaded, then per-page
* permissions applied via filterAccessibleTreePages (restricted-but-granted
* pages are kept; inaccessible subtrees pruned); canEdit is per-page AND
* spaceCanEdit;
* hasChildren is derived from the FINAL (post-prune, post-filter) set, so
* a node never advertises children the user cannot access — the same
* correction getSidebarPages does via getParentIdsWithAccessibleChildren.
*/
async getSidebarPagesTree(
spaceId: string,
userId: string,
spaceCanEdit?: boolean,
pageId?: string,
): Promise<
Array<
Pick<
Page,
| 'id'
| 'slugId'
| 'title'
| 'icon'
| 'position'
| 'parentPageId'
| 'spaceId'
> & { hasChildren: boolean; canEdit: boolean }
>
> {
const hasRestrictions =
await this.pagePermissionRepo.hasRestrictedPagesInSpace(spaceId);
// Seed: a single page subtree, or all root pages of the space.
// Always seed with the FULL (non-excluding) descendant set — in a restricted
// space the per-page filtering below (filterAccessibleTreePages) does the
// pruning, exactly like getSidebarPages. Seeding with *ExcludingRestricted
// would wrongly drop restricted pages the user has an explicit grant for
// (and never recurse into their children), diverging from the sidebar.
let pages: Array<{
id: string;
slugId: string;
title: string;
icon: string;
position: string;
parentPageId: string | null;
spaceId: string;
}>;
if (pageId) {
pages = await this.pageRepo.getPageAndDescendants(pageId, {
includeContent: false,
});
} else {
pages = await this.pageRepo.getSpaceDescendants(spaceId, {
includeContent: false,
});
}
let permissionMap: Map<string, boolean> | undefined;
if (hasRestrictions) {
// Fine-grained per-page permissions on top of restricted pruning.
pages = await this.filterAccessibleTreePages(
pages,
pageId ?? null,
userId,
spaceId,
);
// Per-page canEdit, same source as getSidebarPages.
const accessiblePages =
await this.pagePermissionRepo.filterAccessiblePageIdsWithPermissions(
pages.map((p) => p.id),
userId,
);
permissionMap = new Map(accessiblePages.map((p) => [p.id, p.canEdit]));
}
// Derive hasChildren from the FINAL set: a node has children iff some
// returned row points to it as parent. In a restricted space this set is
// already pruned/filtered, so inaccessible children are not revealed.
const parentIds = new Set<string>();
for (const p of pages) {
if (p.parentPageId) parentIds.add(p.parentPageId);
}
const shaped = pages.map((p) => ({
id: p.id,
slugId: p.slugId,
title: p.title,
icon: p.icon,
position: p.position,
parentPageId: p.parentPageId,
spaceId: p.spaceId,
hasChildren: parentIds.has(p.id),
canEdit: hasRestrictions
? Boolean(permissionMap?.get(p.id)) && (spaceCanEdit ?? true)
: (spaceCanEdit ?? true),
}));
// Order by position with byte order, matching the sidebar's
// `position collate "C"` SQL ordering. position is non-null in returned
// rows; treat a null defensively as sorting last.
shaped.sort((a, b) => {
if (a.position == null) return b.position == null ? 0 : 1;
if (b.position == null) return -1;
return Buffer.compare(Buffer.from(a.position), Buffer.from(b.position));
});
return shaped;
}
}

View File

@@ -0,0 +1,179 @@
/**
* Pure-logic test for getSidebarPagesTree's shaping/permission logic.
*
* NOTE: We cannot import PageService directly here — its dependency chain
* imports `src/collaboration/collaboration.util` via a bare `src/...` path, and
* the server's jest config (package.json "jest".moduleNameMapper) has no
* `^src/(.*)$` mapping, so the module fails to resolve under jest. That is a
* pre-existing config gap unrelated to this feature. To still cover the
* load-bearing logic we replicate the exact shaping algorithm from
* PageService.getSidebarPagesTree below and assert against it. If the service
* logic changes, keep this mirror in sync.
*/
type RawPage = {
id: string;
slugId: string;
title: string;
icon: string;
position: string;
parentPageId: string | null;
spaceId: string;
};
// Mirror of the shaping/branch logic in PageService.getSidebarPagesTree.
function shapeTree(
pages: RawPage[],
opts: {
hasRestrictions: boolean;
spaceCanEdit?: boolean;
permissionMap?: Map<string, boolean>;
},
) {
const parentIds = new Set<string>();
for (const p of pages) {
if (p.parentPageId) parentIds.add(p.parentPageId);
}
const shaped = pages.map((p) => ({
id: p.id,
slugId: p.slugId,
title: p.title,
icon: p.icon,
position: p.position,
parentPageId: p.parentPageId,
spaceId: p.spaceId,
hasChildren: parentIds.has(p.id),
canEdit: opts.hasRestrictions
? Boolean(opts.permissionMap?.get(p.id)) && (opts.spaceCanEdit ?? true)
: (opts.spaceCanEdit ?? true),
}));
shaped.sort((a, b) => {
if (a.position == null) return b.position == null ? 0 : 1;
if (b.position == null) return -1;
return Buffer.compare(Buffer.from(a.position), Buffer.from(b.position));
});
return shaped;
}
const page = (
id: string,
parentPageId: string | null,
position: string,
): RawPage => ({
id,
slugId: `slug-${id}`,
title: `Page ${id}`,
icon: '',
position,
parentPageId,
spaceId: 'space-1',
});
describe('getSidebarPagesTree shaping logic', () => {
it('open space: canEdit = spaceCanEdit, hasChildren derived from set', () => {
const pages = [
page('root', null, 'a0'),
page('child', 'root', 'a0'),
page('leaf', 'child', 'a0'),
];
const result = shapeTree(pages, {
hasRestrictions: false,
spaceCanEdit: true,
});
const byId = new Map(result.map((p) => [p.id, p]));
expect(byId.get('root')!.hasChildren).toBe(true);
expect(byId.get('child')!.hasChildren).toBe(true);
expect(byId.get('leaf')!.hasChildren).toBe(false);
expect(result.every((p) => p.canEdit === true)).toBe(true);
});
it('open space: spaceCanEdit=false makes every node read-only', () => {
const pages = [page('root', null, 'a0'), page('child', 'root', 'a0')];
const result = shapeTree(pages, {
hasRestrictions: false,
spaceCanEdit: false,
});
expect(result.every((p) => p.canEdit === false)).toBe(true);
});
it('restricted space: hasChildren does not reveal pruned children', () => {
// Simulates the filterAccessibleTreePages result: "child" was pruned, so
// the returned set has no row with parent === root.
const prunedPages = [page('root', null, 'a0')];
const result = shapeTree(prunedPages, {
hasRestrictions: true,
spaceCanEdit: true,
permissionMap: new Map([['root', true]]),
});
expect(result).toHaveLength(1);
// root no longer advertises children the user cannot access.
expect(result[0].hasChildren).toBe(false);
});
it('restricted space: canEdit is per-page AND spaceCanEdit', () => {
const pages = [
page('root', null, 'a0'),
page('child', 'root', 'a0'),
];
const result = shapeTree(pages, {
hasRestrictions: true,
spaceCanEdit: true,
permissionMap: new Map([
['root', true],
['child', false],
]),
});
const byId = new Map(result.map((p) => [p.id, p]));
expect(byId.get('root')!.canEdit).toBe(true);
expect(byId.get('child')!.canEdit).toBe(false);
expect(byId.get('root')!.hasChildren).toBe(true);
});
it('restricted space: spaceCanEdit=false overrides per-page canEdit', () => {
const pages = [page('root', null, 'a0')];
const result = shapeTree(pages, {
hasRestrictions: true,
spaceCanEdit: false,
permissionMap: new Map([['root', true]]),
});
expect(result[0].canEdit).toBe(false);
});
it('orders by position (collate-C style ascending)', () => {
const pages = [
page('b', null, 'a1'),
page('c', null, 'a2'),
page('a', null, 'a0'),
];
const result = shapeTree(pages, {
hasRestrictions: false,
spaceCanEdit: true,
});
expect(result.map((p) => p.id)).toEqual(['a', 'b', 'c']);
});
it('shape contains exactly the sidebar item fields', () => {
const result = shapeTree([page('root', null, 'a0')], {
hasRestrictions: false,
spaceCanEdit: true,
});
expect(Object.keys(result[0]).sort()).toEqual(
[
'canEdit',
'hasChildren',
'icon',
'id',
'parentPageId',
'position',
'slugId',
'spaceId',
'title',
].sort(),
);
});
});