Merge remote-tracking branch 'gitea/develop' into HEAD
# Conflicts: # CHANGELOG.md # packages/mcp/build/index.js # packages/mcp/build/lib/auth-utils.js # packages/mcp/build/lib/docmost-schema.js # packages/mcp/build/lib/markdown-converter.js
This commit is contained in:
claude code agent 227
@@ -33,6 +33,11 @@ export class CollaborationGateway {
|
||||
// @ts-ignore
|
||||
private readonly redisSync: RedisSyncExtension<CollabEventHandlers> | null =
|
||||
null;
|
||||
// Source ioredis client that RedisSyncExtension duplicates into its pub/sub
|
||||
// pair. The extension's onDestroy only disconnects those duplicates, so we
|
||||
// keep a reference here and disconnect the source ourselves on shutdown
|
||||
// (otherwise the socket leaks and jest never exits in e2e).
|
||||
private redisClient: RedisClient | null = null;
|
||||
private readonly withRedis: boolean;
|
||||
|
||||
constructor(
|
||||
@@ -57,16 +62,17 @@ export class CollaborationGateway {
|
||||
});
|
||||
|
||||
if (this.withRedis) {
|
||||
this.redisClient = new RedisClient({
|
||||
host: this.redisConfig.host,
|
||||
port: this.redisConfig.port,
|
||||
password: this.redisConfig.password,
|
||||
db: this.redisConfig.db,
|
||||
family: this.redisConfig.family,
|
||||
retryStrategy: createRetryStrategy(),
|
||||
});
|
||||
// @ts-ignore
|
||||
this.redisSync = new RedisSyncExtension({
|
||||
redis: new RedisClient({
|
||||
host: this.redisConfig.host,
|
||||
port: this.redisConfig.port,
|
||||
password: this.redisConfig.password,
|
||||
db: this.redisConfig.db,
|
||||
family: this.redisConfig.family,
|
||||
retryStrategy: createRetryStrategy(),
|
||||
}),
|
||||
redis: this.redisClient,
|
||||
serverId: `collab-${os?.hostname()}-${nanoid(10)}`,
|
||||
prefix: 'collab',
|
||||
pack,
|
||||
@@ -223,5 +229,10 @@ export class CollaborationGateway {
|
||||
});
|
||||
|
||||
await this.hocuspocus.hooks('onDestroy', { instance: this.hocuspocus });
|
||||
|
||||
// RedisSyncExtension.onDestroy (run via the hook above) disconnects only the
|
||||
// duplicated pub/sub clients; the source client created here is ours to close.
|
||||
this.redisClient?.disconnect();
|
||||
this.redisClient = null;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -36,6 +36,7 @@ import {
|
||||
Mention,
|
||||
Subpages,
|
||||
Highlight,
|
||||
Spoiler,
|
||||
Indent,
|
||||
UniqueID,
|
||||
Columns,
|
||||
@@ -82,6 +83,7 @@ export const tiptapExtensions = [
|
||||
Superscript,
|
||||
SubScript,
|
||||
Highlight,
|
||||
Spoiler,
|
||||
Typography,
|
||||
TrailingNode,
|
||||
TextStyle,
|
||||
|
||||
@@ -205,31 +205,203 @@ describe('PersistenceExtension.onStoreDocument — Approach-A boundary snapshot'
|
||||
expect(historyQueue.add).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
// #206 persist-6 — RED (it.failing): a momentarily-empty live Y.Doc must not
|
||||
// overwrite non-empty persisted content. `onStoreDocument` empty-guards the
|
||||
// LOAD path but not the STORE path, so today an empty doc (a client/agent
|
||||
// glitch, a bad merge, an emptying transclusion) is written straight over the
|
||||
// page and the content is wiped silently. A store-side empty-guard is a real
|
||||
// behaviour change (a deliberate "select-all + delete" is also empty), so it
|
||||
// is left UNFIXED pending a product decision; this documents the data-loss
|
||||
// path and flips to a normal passing test the moment the guard lands.
|
||||
it.failing(
|
||||
'does NOT overwrite non-empty content with a momentarily-empty live doc (persist-6)',
|
||||
async () => {
|
||||
const emptyDoc = { type: 'doc', content: [{ type: 'paragraph' }] };
|
||||
const document = ydocFor(emptyDoc);
|
||||
pageRepo.findById.mockResolvedValue({
|
||||
...persistedHumanPage('IGNORED'),
|
||||
content: doc('IMPORTANT RICH CONTENT'),
|
||||
});
|
||||
// #206 persist-6 / #248 — a momentarily-empty live Y.Doc must not overwrite
|
||||
// non-empty persisted content. The store-side empty-guard blocks an empty doc
|
||||
// (a client/agent glitch, a bad merge, an emptying transclusion) from wiping
|
||||
// the page silently when NO intentional-clear signal is present.
|
||||
it('does NOT overwrite non-empty content with a momentarily-empty live doc (persist-6)', async () => {
|
||||
const emptyDoc = { type: 'doc', content: [{ type: 'paragraph' }] };
|
||||
const document = ydocFor(emptyDoc);
|
||||
pageRepo.findById.mockResolvedValue({
|
||||
...persistedHumanPage('IGNORED'),
|
||||
content: doc('IMPORTANT RICH CONTENT'),
|
||||
});
|
||||
|
||||
await ext.onStoreDocument(buildData(document, 'user') as any);
|
||||
await ext.onStoreDocument(buildData(document, 'user') as any);
|
||||
|
||||
// Desired contract: the empty incoming doc is rejected and the rich page
|
||||
// survives. Today updatePage is called with the empty content (data loss).
|
||||
expect(pageRepo.updatePage).not.toHaveBeenCalled();
|
||||
},
|
||||
);
|
||||
// The empty incoming doc is rejected and the rich page survives.
|
||||
expect(pageRepo.updatePage).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
// #248 — an empty-over-empty store is allowed (nothing to lose); the guard
|
||||
// only protects non-empty persisted content.
|
||||
it('allows an empty store over already-empty content (#248)', async () => {
|
||||
const liveEmptyDoc = { type: 'doc', content: [{ type: 'paragraph' }] };
|
||||
const document = ydocFor(liveEmptyDoc);
|
||||
// Stored content is empty per isEmptyParagraphDoc (paragraph with content:[])
|
||||
// but NOT deep-equal to the normalized live doc, so the unchanged
|
||||
// short-circuit is skipped and the empty-guard is genuinely reached.
|
||||
pageRepo.findById.mockResolvedValue({
|
||||
...persistedHumanPage('IGNORED'),
|
||||
content: { type: 'doc', content: [{ type: 'paragraph', content: [] }] },
|
||||
});
|
||||
|
||||
await ext.onStoreDocument(buildData(document, 'user') as any);
|
||||
|
||||
expect(pageRepo.updatePage).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
// #251 — REAL-PATH regression test. The intentional-clear signal is set via
|
||||
// the actual transport seam (ext.onStateless with the exact stateless payload
|
||||
// the client's IntentionalClear extension sends), NOT a hand-injected
|
||||
// context.intentionalClear poke. We then run the debounced store with an empty
|
||||
// live doc over non-empty persisted content and assert the empty write goes
|
||||
// through — i.e. the clear persists.
|
||||
it('persists an intentional clear signalled via the real stateless transport (#251)', async () => {
|
||||
const documentName = `page.${PAGE_ID}`;
|
||||
const emptyDoc = { type: 'doc', content: [{ type: 'paragraph' }] };
|
||||
const document = ydocFor(emptyDoc);
|
||||
pageRepo.findById.mockResolvedValue({
|
||||
...persistedHumanPage('IGNORED'),
|
||||
content: doc('IMPORTANT RICH CONTENT'),
|
||||
});
|
||||
|
||||
// The client signalled a deliberate clear over the live connection.
|
||||
await ext.onStateless({
|
||||
connection: { readOnly: false } as any,
|
||||
documentName,
|
||||
document: document as any,
|
||||
payload: JSON.stringify({ type: 'intentional-clear' }),
|
||||
} as any);
|
||||
|
||||
await ext.onStoreDocument(buildData(document, 'user') as any);
|
||||
|
||||
// The empty doc was written (the clear persisted). The persisted content is
|
||||
// the Y.Doc round-trip of the empty doc (attrs normalized), so compare
|
||||
// against fromYdoc rather than the raw literal.
|
||||
expect(pageRepo.updatePage).toHaveBeenCalledTimes(1);
|
||||
const expectedEmpty = TiptapTransformer.fromYdoc(document, 'default');
|
||||
expect(pageRepo.updatePage.mock.calls[0][0].content).toEqual(expectedEmpty);
|
||||
});
|
||||
|
||||
// #251 — retry correctness: a transient DB failure on the FIRST attempt must
|
||||
// not silently drop the clear. The intentional-clear flag is consumed ONCE
|
||||
// before the retry loop, so when attempt 1's updatePage throws (tx rolls back,
|
||||
// but the in-memory flag delete cannot roll back) the retry on attempt 2 still
|
||||
// sees the clear as allowed and writes the empty doc. On the pre-fix code
|
||||
// (consumeIntentionalClear called INSIDE the loop) attempt 1 consumed the flag,
|
||||
// attempt 2 re-read it as absent and the empty-guard BLOCKED the write — so
|
||||
// updatePage would be called once and the clear would be lost. This test fails
|
||||
// on that ordering and passes after the hoist.
|
||||
it('persists an intentional clear even when the first store attempt fails transiently (#251)', async () => {
|
||||
const documentName = `page.${PAGE_ID}`;
|
||||
const emptyDoc = { type: 'doc', content: [{ type: 'paragraph' }] };
|
||||
const document = ydocFor(emptyDoc);
|
||||
// The page stays non-empty in the DB across both attempts (the rolled-back
|
||||
// first attempt never changed it), exactly the failure scenario the WARNING
|
||||
// describes.
|
||||
pageRepo.findById.mockResolvedValue({
|
||||
...persistedHumanPage('IGNORED'),
|
||||
content: doc('IMPORTANT RICH CONTENT'),
|
||||
});
|
||||
|
||||
let attempts = 0;
|
||||
pageRepo.updatePage.mockImplementation(async () => {
|
||||
attempts += 1;
|
||||
if (attempts === 1) throw new Error('deadlock detected'); // transient
|
||||
callOrder.push('updatePage');
|
||||
});
|
||||
|
||||
// The client signalled a deliberate clear over the live connection.
|
||||
await ext.onStateless({
|
||||
connection: { readOnly: false } as any,
|
||||
documentName,
|
||||
document: document as any,
|
||||
payload: JSON.stringify({ type: 'intentional-clear' }),
|
||||
} as any);
|
||||
|
||||
await ext.onStoreDocument(buildData(document, 'user') as any);
|
||||
|
||||
// First attempt failed and rolled back; the retry still honoured the clear
|
||||
// and wrote the empty doc (the clear survived the retry).
|
||||
expect(pageRepo.updatePage).toHaveBeenCalledTimes(2);
|
||||
const expectedEmpty = TiptapTransformer.fromYdoc(document, 'default');
|
||||
expect(pageRepo.updatePage.mock.calls[1][0].content).toEqual(expectedEmpty);
|
||||
});
|
||||
|
||||
// #251 — the signal is single-use: it is consumed by the first empty store,
|
||||
// so a SECOND accidental empty (no fresh signal) is still blocked.
|
||||
it('consumes the intentional-clear signal once; a later empty is blocked (#251)', async () => {
|
||||
const documentName = `page.${PAGE_ID}`;
|
||||
const emptyDoc = { type: 'doc', content: [{ type: 'paragraph' }] };
|
||||
pageRepo.findById.mockResolvedValue({
|
||||
...persistedHumanPage('IGNORED'),
|
||||
content: doc('IMPORTANT RICH CONTENT'),
|
||||
});
|
||||
|
||||
await ext.onStateless({
|
||||
connection: { readOnly: false } as any,
|
||||
documentName,
|
||||
document: ydocFor(emptyDoc) as any,
|
||||
payload: JSON.stringify({ type: 'intentional-clear' }),
|
||||
} as any);
|
||||
|
||||
// First empty store consumes the signal and writes.
|
||||
await ext.onStoreDocument(buildData(ydocFor(emptyDoc), 'user') as any);
|
||||
expect(pageRepo.updatePage).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Re-arm findById to non-empty (as if content came back) and fire another
|
||||
// empty store WITHOUT a new signal — the guard must block it.
|
||||
pageRepo.updatePage.mockClear();
|
||||
pageRepo.findById.mockResolvedValue({
|
||||
...persistedHumanPage('IGNORED'),
|
||||
content: doc('IMPORTANT RICH CONTENT'),
|
||||
});
|
||||
await ext.onStoreDocument(buildData(ydocFor(emptyDoc), 'user') as any);
|
||||
expect(pageRepo.updatePage).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
// #251 — a read-only connection cannot arm the clear, so its empty store is
|
||||
// still blocked (defends the guard against a read-only spoof).
|
||||
it('ignores an intentional-clear signal from a read-only connection (#251)', async () => {
|
||||
const documentName = `page.${PAGE_ID}`;
|
||||
const emptyDoc = { type: 'doc', content: [{ type: 'paragraph' }] };
|
||||
const document = ydocFor(emptyDoc);
|
||||
pageRepo.findById.mockResolvedValue({
|
||||
...persistedHumanPage('IGNORED'),
|
||||
content: doc('IMPORTANT RICH CONTENT'),
|
||||
});
|
||||
|
||||
await ext.onStateless({
|
||||
connection: { readOnly: true } as any,
|
||||
documentName,
|
||||
document: document as any,
|
||||
payload: JSON.stringify({ type: 'intentional-clear' }),
|
||||
} as any);
|
||||
|
||||
await ext.onStoreDocument(buildData(document, 'user') as any);
|
||||
|
||||
expect(pageRepo.updatePage).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
// #251 — a non-empty store between the signal and the empty store drops the
|
||||
// pending flag ("cleared then retyped" can't leave a usable signal behind).
|
||||
it('drops a pending clear when a non-empty store intervenes (#251)', async () => {
|
||||
const documentName = `page.${PAGE_ID}`;
|
||||
const emptyDoc = { type: 'doc', content: [{ type: 'paragraph' }] };
|
||||
|
||||
await ext.onStateless({
|
||||
connection: { readOnly: false } as any,
|
||||
documentName,
|
||||
document: ydocFor(emptyDoc) as any,
|
||||
payload: JSON.stringify({ type: 'intentional-clear' }),
|
||||
} as any);
|
||||
|
||||
// A non-empty store lands first → consumes/drops the stale flag.
|
||||
pageRepo.findById.mockResolvedValue(persistedHumanPage('NEW HUMAN TEXT'));
|
||||
await ext.onStoreDocument(
|
||||
buildData(ydocFor(doc('NEW HUMAN TEXT')), 'user') as any,
|
||||
);
|
||||
pageRepo.updatePage.mockClear();
|
||||
|
||||
// Now an empty store with no fresh signal must be blocked.
|
||||
pageRepo.findById.mockResolvedValue({
|
||||
...persistedHumanPage('IGNORED'),
|
||||
content: doc('IMPORTANT RICH CONTENT'),
|
||||
});
|
||||
await ext.onStoreDocument(buildData(ydocFor(emptyDoc), 'user') as any);
|
||||
expect(pageRepo.updatePage).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
// persist-1 — when every attempt fails the hook must NOT report a phantom
|
||||
// success: no "page.updated" badge broadcast and no history snapshot for
|
||||
|
||||
@@ -4,6 +4,7 @@ import {
|
||||
onChangePayload,
|
||||
onDisconnectPayload,
|
||||
onLoadDocumentPayload,
|
||||
onStatelessPayload,
|
||||
onStoreDocumentPayload,
|
||||
} from '@hocuspocus/server';
|
||||
import * as Y from 'yjs';
|
||||
@@ -42,6 +43,35 @@ import {
|
||||
} from '../constants';
|
||||
import { TransclusionService } from '../../core/page/transclusion/transclusion.service';
|
||||
|
||||
/**
|
||||
* #251 — wire format of the client→server stateless message that signals a
|
||||
* deliberate page clear. The client (IntentionalClear editor extension) sends
|
||||
* `{ type: INTENTIONAL_CLEAR_MESSAGE_TYPE }`; the document is taken from the
|
||||
* connection, not the payload, so the signal cannot be aimed at another page.
|
||||
*/
|
||||
export const INTENTIONAL_CLEAR_MESSAGE_TYPE = 'intentional-clear';
|
||||
|
||||
/**
|
||||
* #251 — how long an intentional-clear signal stays "pending" before it is
|
||||
* ignored. The signal is set on the clearing keystroke but consumed by the
|
||||
* DEBOUNCED onStoreDocument, so the TTL must comfortably exceed the collab
|
||||
* store debounce window (hocuspocus is configured with maxDebounce = 45s in
|
||||
* collaboration.gateway.ts). 60s leaves a margin while keeping the window for a
|
||||
* stale flag small; on top of the TTL, any non-empty store immediately drops a
|
||||
* pending flag (see onStoreDocument), so a "cleared then retyped" sequence can
|
||||
* never leave a usable flag behind.
|
||||
*
|
||||
* Known fail-safe limitation: the flag lives only in this node's process memory.
|
||||
* If document ownership transfers to another node, or this node crashes/restarts,
|
||||
* between the stateless signal (set on node A) and the debounced store, the
|
||||
* in-memory flag is lost and the clear is silently NOT applied — the store-side
|
||||
* empty-guard then reloads the document non-empty from the DB. This is
|
||||
* deliberately fail-safe (a lost flag preserves content rather than destroying
|
||||
* it), but it is a documented limitation, not a guarantee that every deliberate
|
||||
* clear survives a node handoff.
|
||||
*/
|
||||
export const INTENTIONAL_CLEAR_TTL_MS = 60_000;
|
||||
|
||||
/**
|
||||
* Resolve the provenance source for a coalesced snapshot.
|
||||
*
|
||||
@@ -107,6 +137,13 @@ export class PersistenceExtension implements Extension {
|
||||
// coalescing window" per document and OR it across all edits in the window,
|
||||
// so the snapshot is marked 'agent' regardless of who wrote last.
|
||||
private agentTouched: Map<string, boolean> = new Map();
|
||||
// #251 — per-document "intentional clear pending" flags. Keyed by
|
||||
// documentName, value = expiry timestamp (ms). Set by onStateless when the
|
||||
// client reports a deliberate clear; consumed once by the next
|
||||
// onStoreDocument empty-guard branch. This is the per-EDIT channel the
|
||||
// per-connection context cannot provide (a clear is an edit event, but the
|
||||
// store is debounced and connection context is fixed at authentication).
|
||||
private intentionalClear: Map<string, number> = new Map();
|
||||
|
||||
constructor(
|
||||
private readonly pageRepo: PageRepo,
|
||||
@@ -230,6 +267,19 @@ export class PersistenceExtension implements Extension {
|
||||
this.consumeAgentTouched(documentName),
|
||||
context?.actor,
|
||||
);
|
||||
// #251 — consume the intentional-clear flag ONCE, BEFORE the retry loop
|
||||
// (like consumeContributors / consumeAgentTouched above). consumeIntentional-
|
||||
// Clear ALWAYS deletes the in-memory Map entry, but a tx rollback cannot
|
||||
// un-delete it. Calling it INSIDE the loop meant: a clear armed for attempt 1
|
||||
// was consumed there, attempt 1's updatePage threw a transient error and
|
||||
// rolled back, then attempt 2 re-read non-empty content and saw the flag
|
||||
// already gone — silently downgrading the retry into a BLOCKED write, so the
|
||||
// user's deliberate clear was dropped. Hoisting makes the decision stable
|
||||
// across every attempt. This single call also preserves the "a non-empty
|
||||
// store drops a pending flag" semantics (the cleared-then-retyped case):
|
||||
// every store consumes the flag here regardless of incoming emptiness, so a
|
||||
// subsequent non-empty store can never leave a usable flag behind.
|
||||
const allowIntentionalClear = this.consumeIntentionalClear(documentName);
|
||||
|
||||
// Persist with a small bounded retry. The in-memory Y.Doc is the ONLY copy
|
||||
// of the latest edit until this hook returns: hocuspocus destroys/unloads the
|
||||
@@ -260,6 +310,46 @@ export class PersistenceExtension implements Extension {
|
||||
return;
|
||||
}
|
||||
|
||||
// #206 persist-6 / #248 — store-side empty-guard. A momentarily-empty
|
||||
// live Y.Doc (a client/agent glitch, a bad merge, a transclusion that
|
||||
// emptied) must NOT overwrite non-empty persisted content. The LOAD
|
||||
// path already guards emptiness (onLoadDocument only hydrates from db
|
||||
// when the live doc isEmpty); the STORE path did not, so an empty
|
||||
// serialization was written straight over the page, wiping it
|
||||
// silently.
|
||||
//
|
||||
// #251 — the ONE legitimate empty-over-non-empty write is a user who
|
||||
// deliberately clears the page. That intent arrives out-of-band as a
|
||||
// stateless message, NOT from the doc content, which is why it cannot
|
||||
// be spoofed for non-clear writes: the flag is only ever read on this
|
||||
// empty-incoming branch, so the worst a forged signal can do is clear
|
||||
// a page the connection may already edit. The flag was consumed ONCE
|
||||
// before the retry loop (`allowIntentionalClear`) so the decision is
|
||||
// stable across retries; a non-empty store still drops any pending
|
||||
// flag via that same hoisted consume (a "cleared then retyped"
|
||||
// sequence can't leave a usable one behind).
|
||||
const incomingEmpty = isEmptyParagraphDoc(tiptapJson as any);
|
||||
if (
|
||||
incomingEmpty &&
|
||||
page.content &&
|
||||
!isEmptyParagraphDoc(page.content as any)
|
||||
) {
|
||||
if (allowIntentionalClear) {
|
||||
this.logger.debug(
|
||||
`Intentional clear for ${pageId}: persisting empty doc over ` +
|
||||
`non-empty content (user-signalled)`,
|
||||
);
|
||||
// fall through — the empty write is allowed exactly once.
|
||||
} else {
|
||||
this.logger.warn(
|
||||
`Skipping store for ${pageId}: empty live doc would overwrite ` +
|
||||
`non-empty persisted content`,
|
||||
);
|
||||
page = null;
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
let contributorIds = undefined;
|
||||
try {
|
||||
const existingContributors = page.contributorIds || [];
|
||||
@@ -404,6 +494,37 @@ export class PersistenceExtension implements Extension {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* #251 — receive the client's deliberate-clear signal. Records a short-lived,
|
||||
* single-use pending flag for the originating document so the next
|
||||
* onStoreDocument may let one empty-over-non-empty write through the guard.
|
||||
*
|
||||
* Hardening: read-only connections cannot arm the flag, and the document is
|
||||
* taken from the connection (`data.documentName`), never the payload, so a
|
||||
* client cannot target a page it isn't editing. The flag only ever RELAXES
|
||||
* the guard for an empty write (a clear); it can never force or alter a
|
||||
* non-empty write, so it is not a guard bypass for normal content.
|
||||
*/
|
||||
async onStateless(data: onStatelessPayload) {
|
||||
const { connection, documentName, payload } = data;
|
||||
|
||||
if (connection?.readOnly) return;
|
||||
|
||||
let message: { type?: string } | undefined;
|
||||
try {
|
||||
message = JSON.parse(payload);
|
||||
} catch {
|
||||
return; // unrelated / malformed stateless message
|
||||
}
|
||||
|
||||
if (message?.type !== INTENTIONAL_CLEAR_MESSAGE_TYPE) return;
|
||||
|
||||
this.intentionalClear.set(
|
||||
documentName,
|
||||
Date.now() + INTENTIONAL_CLEAR_TTL_MS,
|
||||
);
|
||||
}
|
||||
|
||||
async onChange(data: onChangePayload) {
|
||||
const documentName = data.documentName;
|
||||
const userId = data.context?.user?.id;
|
||||
@@ -427,6 +548,7 @@ export class PersistenceExtension implements Extension {
|
||||
const documentName = data.documentName;
|
||||
this.contributors.delete(documentName);
|
||||
this.agentTouched.delete(documentName);
|
||||
this.intentionalClear.delete(documentName);
|
||||
}
|
||||
|
||||
private consumeContributors(documentName: string): string[] {
|
||||
@@ -444,6 +566,18 @@ export class PersistenceExtension implements Extension {
|
||||
return touched;
|
||||
}
|
||||
|
||||
/**
|
||||
* #251 — read and clear the intentional-clear flag for this document. Returns
|
||||
* true only if a flag was pending AND still within its TTL. Always deletes the
|
||||
* entry so the signal is strictly single-use (one clear → one allowed empty
|
||||
* write); an expired flag is treated as absent (guard still blocks).
|
||||
*/
|
||||
private consumeIntentionalClear(documentName: string): boolean {
|
||||
const expiry = this.intentionalClear.get(documentName);
|
||||
this.intentionalClear.delete(documentName);
|
||||
return expiry !== undefined && Date.now() < expiry;
|
||||
}
|
||||
|
||||
private async enqueuePageHistory(
|
||||
page: Page,
|
||||
lastUpdatedSource: string,
|
||||
|
||||
@@ -0,0 +1,278 @@
|
||||
import * as Y from 'yjs';
|
||||
import { getSchema } from '@tiptap/core';
|
||||
import {
|
||||
initProseMirrorDoc,
|
||||
absolutePositionToRelativePosition,
|
||||
prosemirrorJSONToYDoc,
|
||||
} from '@tiptap/y-tiptap';
|
||||
import { tiptapExtensions } from './collaboration.util';
|
||||
import {
|
||||
setYjsMark,
|
||||
removeYjsMarkByAttribute,
|
||||
updateYjsMarkAttribute,
|
||||
type YjsSelection,
|
||||
} from './yjs.util';
|
||||
|
||||
/**
|
||||
* Unit tests for the server-side Yjs mark helpers used by the collaboration
|
||||
* handler to set/resolve/delete comment marks directly on the shared Y.Doc
|
||||
* (collaboration.handler.ts: setCommentMark / resolveCommentMark).
|
||||
*
|
||||
* The fragment shape mirrors production exactly: a `default` XmlFragment whose
|
||||
* children are block XmlElements (paragraph) holding XmlText runs. For setYjsMark
|
||||
* the selection is a pair of Yjs RelativePosition JSONs (what the client sends);
|
||||
* we synthesize them from known ProseMirror absolute positions via
|
||||
* absolutePositionToRelativePosition so the marked range is deterministic.
|
||||
*/
|
||||
|
||||
const schema = getSchema(tiptapExtensions);
|
||||
|
||||
// Build a real Y.Doc from ProseMirror JSON (same path the collab handler uses
|
||||
// via TiptapTransformer) and return the doc + its `default` fragment.
|
||||
function buildFromPm(pmJson: unknown) {
|
||||
const ydoc = prosemirrorJSONToYDoc(
|
||||
schema,
|
||||
pmJson as never,
|
||||
'default',
|
||||
) as unknown as Y.Doc;
|
||||
const fragment = ydoc.getXmlFragment('default');
|
||||
return { ydoc, fragment };
|
||||
}
|
||||
|
||||
// Make a YjsSelection (anchor/head RelativePosition JSON) for two ProseMirror
|
||||
// absolute positions in `fragment`.
|
||||
function selectionFor(
|
||||
fragment: Y.XmlFragment,
|
||||
anchorPos: number,
|
||||
headPos: number,
|
||||
): YjsSelection {
|
||||
const { mapping } = initProseMirrorDoc(fragment, schema);
|
||||
const anchor = absolutePositionToRelativePosition(
|
||||
anchorPos,
|
||||
fragment as never,
|
||||
mapping,
|
||||
);
|
||||
const head = absolutePositionToRelativePosition(
|
||||
headPos,
|
||||
fragment as never,
|
||||
mapping,
|
||||
);
|
||||
return {
|
||||
anchor: Y.relativePositionToJSON(anchor),
|
||||
head: Y.relativePositionToJSON(head),
|
||||
};
|
||||
}
|
||||
|
||||
// The XmlText run of the i-th top-level paragraph.
|
||||
function paragraphText(fragment: Y.XmlFragment, index = 0): Y.XmlText {
|
||||
const para = fragment.get(index) as Y.XmlElement;
|
||||
return para.get(0) as Y.XmlText;
|
||||
}
|
||||
|
||||
// --- raw fragment builder for the remove/update tests (no schema needed) ---
|
||||
//
|
||||
// removeYjsMarkByAttribute / updateYjsMarkAttribute only read item.toDelta() and
|
||||
// call item.format(); they never touch the ProseMirror schema. Build the runs
|
||||
// directly so we control which segment carries which comment attrs.
|
||||
function buildWithComments(
|
||||
segments: Array<{
|
||||
text: string;
|
||||
comment?: { commentId: string; resolved: boolean };
|
||||
}>,
|
||||
): { fragment: Y.XmlFragment; text: Y.XmlText } {
|
||||
const ydoc = new Y.Doc();
|
||||
const fragment = ydoc.getXmlFragment('default');
|
||||
const para = new Y.XmlElement('paragraph');
|
||||
fragment.insert(0, [para]);
|
||||
const text = new Y.XmlText();
|
||||
para.insert(0, [text]);
|
||||
let offset = 0;
|
||||
for (const seg of segments) {
|
||||
text.insert(offset, seg.text);
|
||||
if (seg.comment) {
|
||||
text.format(offset, seg.text.length, { comment: seg.comment });
|
||||
}
|
||||
offset += seg.text.length;
|
||||
}
|
||||
return { fragment, text };
|
||||
}
|
||||
|
||||
describe('setYjsMark', () => {
|
||||
it('applies the mark over exactly the selected sub-range (PM pos 1..6 = "Hello")', () => {
|
||||
const { ydoc, fragment } = buildFromPm({
|
||||
type: 'doc',
|
||||
content: [
|
||||
{ type: 'paragraph', content: [{ type: 'text', text: 'Hello world' }] },
|
||||
],
|
||||
});
|
||||
// PM pos 1 = start of the paragraph text; pos 6 = just after "Hello".
|
||||
const sel = selectionFor(fragment, 1, 6);
|
||||
|
||||
setYjsMark(ydoc as never, fragment, sel, 'comment', {
|
||||
commentId: 'c1',
|
||||
resolved: false,
|
||||
});
|
||||
|
||||
// The run splits: "Hello" carries the comment mark, " world" stays clean.
|
||||
expect(paragraphText(fragment).toDelta()).toEqual([
|
||||
{
|
||||
insert: 'Hello',
|
||||
attributes: { comment: { commentId: 'c1', resolved: false } },
|
||||
},
|
||||
{ insert: ' world' },
|
||||
]);
|
||||
});
|
||||
|
||||
it('normalizes a reversed selection (head before anchor) to the same range', () => {
|
||||
const { ydoc, fragment } = buildFromPm({
|
||||
type: 'doc',
|
||||
content: [
|
||||
{ type: 'paragraph', content: [{ type: 'text', text: 'Hello world' }] },
|
||||
],
|
||||
});
|
||||
// anchor=6, head=1 — reversed; setYjsMark takes min/max so it marks "Hello".
|
||||
const sel = selectionFor(fragment, 6, 1);
|
||||
|
||||
setYjsMark(ydoc as never, fragment, sel, 'comment', {
|
||||
commentId: 'c2',
|
||||
resolved: false,
|
||||
});
|
||||
|
||||
expect(paragraphText(fragment).toDelta()).toEqual([
|
||||
{
|
||||
insert: 'Hello',
|
||||
attributes: { comment: { commentId: 'c2', resolved: false } },
|
||||
},
|
||||
{ insert: ' world' },
|
||||
]);
|
||||
});
|
||||
|
||||
it('marks across two paragraphs (range spans an element boundary)', () => {
|
||||
const { ydoc, fragment } = buildFromPm({
|
||||
type: 'doc',
|
||||
content: [
|
||||
{ type: 'paragraph', content: [{ type: 'text', text: 'aaa' }] },
|
||||
{ type: 'paragraph', content: [{ type: 'text', text: 'bbb' }] },
|
||||
],
|
||||
});
|
||||
// PM positions: "aaa" = 1..4; the </p><p> boundary consumes pos 4 and 5, so
|
||||
// "bbb" starts at pos 6 (chars at 6,7,8). Select pos 2 (inside "aaa") to pos
|
||||
// 8 (after the second "b").
|
||||
const sel = selectionFor(fragment, 2, 8);
|
||||
|
||||
setYjsMark(ydoc as never, fragment, sel, 'comment', {
|
||||
commentId: 'c3',
|
||||
resolved: false,
|
||||
});
|
||||
|
||||
// First paragraph: "a" clean, "aa" marked.
|
||||
expect(paragraphText(fragment, 0).toDelta()).toEqual([
|
||||
{ insert: 'a' },
|
||||
{
|
||||
insert: 'aa',
|
||||
attributes: { comment: { commentId: 'c3', resolved: false } },
|
||||
},
|
||||
]);
|
||||
// Second paragraph: "bb" marked, "b" clean.
|
||||
expect(paragraphText(fragment, 1).toDelta()).toEqual([
|
||||
{
|
||||
insert: 'bb',
|
||||
attributes: { comment: { commentId: 'c3', resolved: false } },
|
||||
},
|
||||
{ insert: 'b' },
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('removeYjsMarkByAttribute', () => {
|
||||
it('removes only the run whose attribute value matches, leaving others', () => {
|
||||
const { fragment, text } = buildWithComments([
|
||||
{ text: 'AAA', comment: { commentId: 'c1', resolved: false } },
|
||||
{ text: 'BBB', comment: { commentId: 'c2', resolved: false } },
|
||||
]);
|
||||
|
||||
removeYjsMarkByAttribute(fragment, 'comment', 'commentId', 'c1');
|
||||
|
||||
// c1's run loses the mark; c2's run is untouched.
|
||||
expect(text.toDelta()).toEqual([
|
||||
{ insert: 'AAA' },
|
||||
{
|
||||
insert: 'BBB',
|
||||
attributes: { comment: { commentId: 'c2', resolved: false } },
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
it('does nothing when no run carries the requested value (no-match branch)', () => {
|
||||
const { fragment, text } = buildWithComments([
|
||||
{ text: 'AAA', comment: { commentId: 'c1', resolved: false } },
|
||||
]);
|
||||
const before = text.toDelta();
|
||||
|
||||
removeYjsMarkByAttribute(fragment, 'comment', 'commentId', 'does-not-exist');
|
||||
|
||||
expect(text.toDelta()).toEqual(before);
|
||||
});
|
||||
|
||||
it('leaves a different mark type alone', () => {
|
||||
// A run carrying only `bold` must survive a comment removal pass.
|
||||
const ydoc = new Y.Doc();
|
||||
const fragment = ydoc.getXmlFragment('default');
|
||||
const para = new Y.XmlElement('paragraph');
|
||||
fragment.insert(0, [para]);
|
||||
const text = new Y.XmlText();
|
||||
para.insert(0, [text]);
|
||||
text.insert(0, 'XYZ');
|
||||
text.format(0, 3, { bold: true });
|
||||
|
||||
removeYjsMarkByAttribute(fragment, 'comment', 'commentId', 'c1');
|
||||
|
||||
expect(text.toDelta()).toEqual([
|
||||
{ insert: 'XYZ', attributes: { bold: true } },
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('updateYjsMarkAttribute', () => {
|
||||
it('merges new attributes into the matching run, preserving the rest', () => {
|
||||
const { fragment, text } = buildWithComments([
|
||||
{ text: 'AAA', comment: { commentId: 'c1', resolved: false } },
|
||||
{ text: 'BBB', comment: { commentId: 'c2', resolved: false } },
|
||||
]);
|
||||
|
||||
updateYjsMarkAttribute(
|
||||
fragment,
|
||||
'comment',
|
||||
{ name: 'commentId', value: 'c1' },
|
||||
{ resolved: true },
|
||||
);
|
||||
|
||||
// c1's run flips resolved=true (commentId preserved via merge); c2 untouched.
|
||||
expect(text.toDelta()).toEqual([
|
||||
{
|
||||
insert: 'AAA',
|
||||
attributes: { comment: { commentId: 'c1', resolved: true } },
|
||||
},
|
||||
{
|
||||
insert: 'BBB',
|
||||
attributes: { comment: { commentId: 'c2', resolved: false } },
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
it('does nothing when no run matches (no-match branch)', () => {
|
||||
const { fragment, text } = buildWithComments([
|
||||
{ text: 'AAA', comment: { commentId: 'c1', resolved: false } },
|
||||
]);
|
||||
const before = text.toDelta();
|
||||
|
||||
updateYjsMarkAttribute(
|
||||
fragment,
|
||||
'comment',
|
||||
{ name: 'commentId', value: 'nope' },
|
||||
{ resolved: true },
|
||||
);
|
||||
|
||||
expect(text.toDelta()).toEqual(before);
|
||||
});
|
||||
});
|
||||
@@ -3,6 +3,8 @@ import { PageRepo } from '@docmost/db/repos/page/page.repo';
|
||||
import { PageEmbeddingRepo } from '@docmost/db/repos/ai-chat/page-embedding.repo';
|
||||
import { KyselyDB } from '@docmost/db/types/kysely.types';
|
||||
import { AiService } from '../../../integrations/ai/ai.service';
|
||||
import { EmbeddingReindexProgressService } from '../../../integrations/ai/embedding-reindex-progress.service';
|
||||
import { AiEmbeddingNotConfiguredException } from '../../../integrations/ai/ai-embedding-not-configured.exception';
|
||||
|
||||
/**
|
||||
* Unit tests for EmbeddingIndexerService.reindexWorkspace's batch control flow.
|
||||
@@ -12,7 +14,8 @@ import { AiService } from '../../../integrations/ai/ai.service';
|
||||
* reindexWorkspace actually touches:
|
||||
* - aiService.getEmbeddingModel -> a model string so the up-front configured
|
||||
* check passes,
|
||||
* - pageRepo.getIdsByWorkspace -> three page ids,
|
||||
* - pageRepo.getEmbeddablePageIds -> three page ids (the embeddable set the
|
||||
* reindex iterates),
|
||||
* - service.reindexPage -> spied per test to drive the per-page outcome.
|
||||
*
|
||||
* The point under test is the catch block: a FATAL provider error (auth/billing)
|
||||
@@ -24,21 +27,30 @@ describe('EmbeddingIndexerService.reindexWorkspace fail-fast', () => {
|
||||
|
||||
function makeService() {
|
||||
const pageRepo = {
|
||||
getIdsByWorkspace: jest.fn().mockResolvedValue(['p1', 'p2', 'p3']),
|
||||
getEmbeddablePageIds: jest.fn().mockResolvedValue(['p1', 'p2', 'p3']),
|
||||
};
|
||||
const pageEmbeddingRepo = {};
|
||||
const aiService = {
|
||||
getEmbeddingModel: jest.fn().mockResolvedValue('some-model'),
|
||||
};
|
||||
// Progress is a best-effort cosmetic store; mock its async methods so the
|
||||
// batch control flow can be tested without Redis.
|
||||
const reindexProgress = {
|
||||
start: jest.fn().mockResolvedValue(undefined),
|
||||
increment: jest.fn().mockResolvedValue(undefined),
|
||||
clear: jest.fn().mockResolvedValue(undefined),
|
||||
get: jest.fn().mockResolvedValue(null),
|
||||
};
|
||||
const db = {};
|
||||
|
||||
const service = new EmbeddingIndexerService(
|
||||
pageRepo as unknown as PageRepo,
|
||||
pageEmbeddingRepo as unknown as PageEmbeddingRepo,
|
||||
aiService as unknown as AiService,
|
||||
reindexProgress as unknown as EmbeddingReindexProgressService,
|
||||
db as unknown as KyselyDB,
|
||||
);
|
||||
return { service, pageRepo, aiService };
|
||||
return { service, pageRepo, aiService, reindexProgress };
|
||||
}
|
||||
|
||||
it('aborts after the first page on a FATAL (401) provider error', async () => {
|
||||
@@ -78,3 +90,100 @@ describe('EmbeddingIndexerService.reindexWorkspace fail-fast', () => {
|
||||
expect(reindexPage).toHaveBeenCalledTimes(3);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* Live reindex-progress reporting: reindexWorkspace must publish a per-workspace
|
||||
* progress record (total at start, done incremented per processed page) and ALWAYS
|
||||
* clear it in a finally — including on a fatal abort and an unconfigured early
|
||||
* return — so the settings status can show the counter climb without ever getting
|
||||
* stuck in a "reindexing" state.
|
||||
*/
|
||||
describe('EmbeddingIndexerService.reindexWorkspace progress', () => {
|
||||
const WORKSPACE_ID = 'ws-1';
|
||||
|
||||
function makeService(pageIds: string[] = ['p1', 'p2', 'p3']) {
|
||||
const pageRepo = {
|
||||
getEmbeddablePageIds: jest.fn().mockResolvedValue(pageIds),
|
||||
};
|
||||
const pageEmbeddingRepo = {};
|
||||
const aiService = {
|
||||
getEmbeddingModel: jest.fn().mockResolvedValue('some-model'),
|
||||
};
|
||||
const reindexProgress = {
|
||||
start: jest.fn().mockResolvedValue(undefined),
|
||||
increment: jest.fn().mockResolvedValue(undefined),
|
||||
clear: jest.fn().mockResolvedValue(undefined),
|
||||
get: jest.fn().mockResolvedValue(null),
|
||||
};
|
||||
const db = {};
|
||||
const service = new EmbeddingIndexerService(
|
||||
pageRepo as unknown as PageRepo,
|
||||
pageEmbeddingRepo as unknown as PageEmbeddingRepo,
|
||||
aiService as unknown as AiService,
|
||||
reindexProgress as unknown as EmbeddingReindexProgressService,
|
||||
db as unknown as KyselyDB,
|
||||
);
|
||||
return { service, pageRepo, aiService, reindexProgress };
|
||||
}
|
||||
|
||||
it('sets total at start, increments done per page, and clears in finally', async () => {
|
||||
const { service, reindexProgress } = makeService(['p1', 'p2', 'p3']);
|
||||
jest.spyOn(service, 'reindexPage').mockResolvedValue(undefined);
|
||||
|
||||
await service.reindexWorkspace(WORKSPACE_ID);
|
||||
|
||||
expect(reindexProgress.start).toHaveBeenCalledWith(WORKSPACE_ID, 3);
|
||||
// One increment per processed page.
|
||||
expect(reindexProgress.increment).toHaveBeenCalledTimes(3);
|
||||
expect(reindexProgress.increment).toHaveBeenCalledWith(WORKSPACE_ID);
|
||||
// Cleared exactly once on completion.
|
||||
expect(reindexProgress.clear).toHaveBeenCalledTimes(1);
|
||||
expect(reindexProgress.clear).toHaveBeenCalledWith(WORKSPACE_ID);
|
||||
});
|
||||
|
||||
it('counts a handled (non-fatal) per-page failure as processed', async () => {
|
||||
const { service, reindexProgress } = makeService(['p1', 'p2', 'p3']);
|
||||
// No statusCode -> non-fatal -> isolate and continue; each counts as done.
|
||||
jest.spyOn(service, 'reindexPage').mockRejectedValue(new Error('boom'));
|
||||
|
||||
await service.reindexWorkspace(WORKSPACE_ID);
|
||||
|
||||
expect(reindexProgress.increment).toHaveBeenCalledTimes(3);
|
||||
expect(reindexProgress.clear).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('clears progress in finally even when a FATAL provider error aborts the batch', async () => {
|
||||
const { service, reindexProgress } = makeService(['p1', 'p2', 'p3']);
|
||||
// A 401 aborts on the first page (re-thrown) — the finally must still clear.
|
||||
jest
|
||||
.spyOn(service, 'reindexPage')
|
||||
.mockRejectedValue({ statusCode: 401, message: 'User not found' });
|
||||
|
||||
await expect(service.reindexWorkspace(WORKSPACE_ID)).rejects.toMatchObject({
|
||||
statusCode: 401,
|
||||
});
|
||||
|
||||
expect(reindexProgress.start).toHaveBeenCalledWith(WORKSPACE_ID, 3);
|
||||
// Aborted page is NOT counted as processed.
|
||||
expect(reindexProgress.increment).not.toHaveBeenCalled();
|
||||
// But progress is still cleared so the run never gets stuck.
|
||||
expect(reindexProgress.clear).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('clears the enqueue-seeded progress on an unconfigured early return', async () => {
|
||||
const { service, aiService, reindexProgress } = makeService();
|
||||
// Embeddings not configured: reindexWorkspace returns early WITHOUT starting
|
||||
// a fresh record, but the finally must still clear the enqueue-time seed.
|
||||
aiService.getEmbeddingModel = jest
|
||||
.fn()
|
||||
.mockRejectedValue(new AiEmbeddingNotConfiguredException());
|
||||
|
||||
await expect(
|
||||
service.reindexWorkspace(WORKSPACE_ID),
|
||||
).resolves.toBeUndefined();
|
||||
|
||||
expect(reindexProgress.start).not.toHaveBeenCalled();
|
||||
expect(reindexProgress.clear).toHaveBeenCalledTimes(1);
|
||||
expect(reindexProgress.clear).toHaveBeenCalledWith(WORKSPACE_ID);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -9,6 +9,7 @@ import { KyselyDB } from '@docmost/db/types/kysely.types';
|
||||
import { InjectKysely } from 'nestjs-kysely';
|
||||
import { executeTx } from '@docmost/db/utils';
|
||||
import { AiService } from '../../../integrations/ai/ai.service';
|
||||
import { EmbeddingReindexProgressService } from '../../../integrations/ai/embedding-reindex-progress.service';
|
||||
import { AiEmbeddingNotConfiguredException } from '../../../integrations/ai/ai-embedding-not-configured.exception';
|
||||
import {
|
||||
describeProviderError,
|
||||
@@ -48,6 +49,7 @@ export class EmbeddingIndexerService {
|
||||
private readonly pageRepo: PageRepo,
|
||||
private readonly pageEmbeddingRepo: PageEmbeddingRepo,
|
||||
private readonly aiService: AiService,
|
||||
private readonly reindexProgress: EmbeddingReindexProgressService,
|
||||
@InjectKysely() private readonly db: KyselyDB,
|
||||
) {}
|
||||
|
||||
@@ -183,7 +185,19 @@ export class EmbeddingIndexerService {
|
||||
}
|
||||
|
||||
/**
|
||||
* (Re)build embeddings for EVERY non-deleted page in a workspace. Used by the
|
||||
* (Re)build embeddings for the EMBEDDABLE page set of a workspace — the same
|
||||
* set countEmbeddablePages counts (via getEmbeddablePageIds): non-deleted pages
|
||||
* that qualify under any of the three clauses of `embeddablePredicate` —
|
||||
* non-empty textContent, OR an empty/null textContent whose ProseMirror
|
||||
* `content` JSON has at least one text node (`"type":"text"`) that `jsonToText`
|
||||
* can extract, OR an already-stored (non-deleted) embedding row — NOT every
|
||||
* non-deleted page. Iterating this set keeps the live `total` equal to the
|
||||
* steady-state denominator, so the progress counter climbs 0 -> total and
|
||||
* matches the before/after DB coverage exactly. A page with truly no
|
||||
* extractable text (empty textContent AND content with only non-text/atom
|
||||
* nodes such as math) is correctly skipped (reindexPage no-ops on it); a page
|
||||
* that lost its text but still has stale embeddings stays in the set (the
|
||||
* EXISTS clause) so it is visited and its stale rows are cleared. Used by the
|
||||
* bulk reindex (WORKSPACE_CREATE_EMBEDDINGS, fired when AI Search is enabled
|
||||
* and by the manual "Reindex now" action).
|
||||
*
|
||||
@@ -194,69 +208,99 @@ export class EmbeddingIndexerService {
|
||||
* the batch.
|
||||
*/
|
||||
async reindexWorkspace(workspaceId: string): Promise<void> {
|
||||
// The whole run is wrapped so the per-workspace progress record is ALWAYS
|
||||
// cleared in the finally — on success, on a fatal-provider abort, on an
|
||||
// unconfigured early-return, or on any unexpected throw — so a failed run
|
||||
// never leaves a stuck "reindexing" state (the status then falls back to the
|
||||
// steady-state DB coverage count). A placeholder record may already exist
|
||||
// (seeded at enqueue time); the finally cleans that too.
|
||||
try {
|
||||
await this.aiService.getEmbeddingModel(workspaceId);
|
||||
} catch (err) {
|
||||
if (err instanceof AiEmbeddingNotConfiguredException) {
|
||||
this.logger.log(
|
||||
`reindexWorkspace: embeddings not configured for workspace ${workspaceId}, skipping`,
|
||||
);
|
||||
return;
|
||||
}
|
||||
throw err;
|
||||
}
|
||||
|
||||
const pageIds = await this.pageRepo.getIdsByWorkspace(workspaceId);
|
||||
const total = pageIds.length;
|
||||
const startedAt = Date.now();
|
||||
this.logger.log(
|
||||
`reindexWorkspace: starting reindex of ${total} page(s) for workspace ${workspaceId}`,
|
||||
);
|
||||
|
||||
let failed = 0;
|
||||
for (let i = 0; i < total; i++) {
|
||||
const pageId = pageIds[i];
|
||||
const position = i + 1;
|
||||
// Log BEFORE the await: if the embedding call hangs, this is the last line
|
||||
// in the log and it names the exact page that is stuck.
|
||||
this.logger.log(
|
||||
`reindexWorkspace: [${position}/${total}] indexing page ${pageId} (workspace ${workspaceId})`,
|
||||
);
|
||||
const pageStartedAt = Date.now();
|
||||
try {
|
||||
await this.reindexPage(pageId);
|
||||
const elapsed = Date.now() - pageStartedAt;
|
||||
if (elapsed >= SLOW_PAGE_MS) {
|
||||
this.logger.warn(
|
||||
`reindexWorkspace: [${position}/${total}] page ${pageId} took ${elapsed}ms`,
|
||||
);
|
||||
}
|
||||
await this.aiService.getEmbeddingModel(workspaceId);
|
||||
} catch (err) {
|
||||
// A fatal provider error (invalid/missing key, no credits) recurs
|
||||
// identically on EVERY remaining page. Abort the whole batch instead of
|
||||
// issuing hundreds of doomed requests against the provider.
|
||||
if (isFatalProviderError(err)) {
|
||||
this.logger.error(
|
||||
`reindexWorkspace: aborting at [${position}/${total}] for workspace ` +
|
||||
`${workspaceId} — fatal provider error, remaining pages would fail ` +
|
||||
`identically: ${describeProviderError(err)}`,
|
||||
if (err instanceof AiEmbeddingNotConfiguredException) {
|
||||
this.logger.log(
|
||||
`reindexWorkspace: embeddings not configured for workspace ${workspaceId}, skipping`,
|
||||
);
|
||||
throw err;
|
||||
return;
|
||||
}
|
||||
// Per-page isolation: one non-fatal failure (incl. an embedding timeout)
|
||||
// must not abort the whole batch.
|
||||
failed++;
|
||||
this.logger.error(
|
||||
`reindexWorkspace: [${position}/${total}] failed to reindex page ${pageId} ` +
|
||||
`after ${Date.now() - pageStartedAt}ms: ${describeProviderError(err)}`,
|
||||
);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
this.logger.log(
|
||||
`reindexWorkspace: done for workspace ${workspaceId}: ` +
|
||||
`${total - failed}/${total} indexed, ${failed} failed in ${Date.now() - startedAt}ms`,
|
||||
);
|
||||
// Iterate the EMBEDDABLE set (same three-clause predicate as
|
||||
// countEmbeddablePages), NOT every non-deleted page: this makes `total`
|
||||
// here equal the steady-state denominator, so the live counter climbs
|
||||
// 0 -> total and matches the before/after DB count exactly (no
|
||||
// 478 -> 500 -> 478 denominator jump). Pages whose text lives in the
|
||||
// ProseMirror `content` JSON (a text node) even with empty text_content ARE
|
||||
// in this set (the content-JSON clause) and get embedded; a page with no
|
||||
// extractable text at all is correctly skipped — reindexPage no-ops on it —
|
||||
// and a page that lost its text but still has stale embeddings IS in this
|
||||
// set (the EXISTS clause) so it is still visited and its stale rows cleared.
|
||||
const pageIds = await this.pageRepo.getEmbeddablePageIds(workspaceId);
|
||||
const total = pageIds.length;
|
||||
const startedAt = Date.now();
|
||||
// Publish the live run progress over this same set (done reset to 0). The
|
||||
// counter increments once per iterated page and reaches exactly `total`,
|
||||
// which equals countEmbeddablePages — the steady-state denominator.
|
||||
await this.reindexProgress.start(workspaceId, total);
|
||||
this.logger.log(
|
||||
`reindexWorkspace: starting reindex of ${total} page(s) for workspace ${workspaceId}`,
|
||||
);
|
||||
|
||||
let failed = 0;
|
||||
for (let i = 0; i < total; i++) {
|
||||
const pageId = pageIds[i];
|
||||
const position = i + 1;
|
||||
// Log BEFORE the await: if the embedding call hangs, this is the last line
|
||||
// in the log and it names the exact page that is stuck.
|
||||
this.logger.log(
|
||||
`reindexWorkspace: [${position}/${total}] indexing page ${pageId} (workspace ${workspaceId})`,
|
||||
);
|
||||
const pageStartedAt = Date.now();
|
||||
try {
|
||||
await this.reindexPage(pageId);
|
||||
// Count this page as processed (matches the [position/total] log).
|
||||
await this.reindexProgress.increment(workspaceId);
|
||||
const elapsed = Date.now() - pageStartedAt;
|
||||
if (elapsed >= SLOW_PAGE_MS) {
|
||||
this.logger.warn(
|
||||
`reindexWorkspace: [${position}/${total}] page ${pageId} took ${elapsed}ms`,
|
||||
);
|
||||
}
|
||||
} catch (err) {
|
||||
// A fatal provider error (invalid/missing key, no credits) recurs
|
||||
// identically on EVERY remaining page. Abort the whole batch instead of
|
||||
// issuing hundreds of doomed requests against the provider. Do NOT count
|
||||
// it as processed — the run aborts here (the finally clears progress).
|
||||
if (isFatalProviderError(err)) {
|
||||
this.logger.error(
|
||||
`reindexWorkspace: aborting at [${position}/${total}] for workspace ` +
|
||||
`${workspaceId} — fatal provider error, remaining pages would fail ` +
|
||||
`identically: ${describeProviderError(err)}`,
|
||||
);
|
||||
throw err;
|
||||
}
|
||||
// Per-page isolation: one non-fatal failure (incl. an embedding timeout)
|
||||
// must not abort the whole batch. A handled failure still advances the
|
||||
// counter (matches the [position/total] log, so done reaches total).
|
||||
failed++;
|
||||
await this.reindexProgress.increment(workspaceId);
|
||||
this.logger.error(
|
||||
`reindexWorkspace: [${position}/${total}] failed to reindex page ${pageId} ` +
|
||||
`after ${Date.now() - pageStartedAt}ms: ${describeProviderError(err)}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
this.logger.log(
|
||||
`reindexWorkspace: done for workspace ${workspaceId}: ` +
|
||||
`${total - failed}/${total} indexed, ${failed} failed in ${Date.now() - startedAt}ms`,
|
||||
);
|
||||
} finally {
|
||||
// Always remove the progress record so the status reverts to the DB count.
|
||||
await this.reindexProgress.clear(workspaceId);
|
||||
}
|
||||
}
|
||||
|
||||
/** Purge ALL embeddings for a workspace (WORKSPACE_DELETE_EMBEDDINGS). */
|
||||
|
||||
@@ -0,0 +1,166 @@
|
||||
import { McpClientsService } from './mcp-clients.service';
|
||||
|
||||
/**
|
||||
* Unit tests for the two security-critical surfaces of McpClientsService that the
|
||||
* sibling specs (ssrf-guard / validate-resolved-addresses / lease) do NOT cover:
|
||||
*
|
||||
* 1. `decryptHeaders` (private) — FAIL-OPEN behavior. A decrypt/parse failure
|
||||
* (e.g. APP_SECRET rotated, tampered blob) must NEVER throw and must NEVER
|
||||
* log the blob: it returns `undefined` so the connect proceeds WITHOUT the
|
||||
* now-unreadable auth headers (which then 401s and the server is skipped),
|
||||
* rather than crashing the whole turn.
|
||||
*
|
||||
* 2. `this.guardedFetch` (private, bound to the SSRF-pinned dispatcher) — the
|
||||
* per-request DNS-rebinding guard. A blocked host (private/loopback/metadata
|
||||
* IP literal, or an unparseable URL) must REJECT before any socket is opened;
|
||||
* a public host is allowed through to the real `fetch` with the pinned
|
||||
* dispatcher attached.
|
||||
*
|
||||
* No network and no DB: the repo + secretBox deps are stubbed, and global `fetch`
|
||||
* is mocked for the single allow-path assertion.
|
||||
*/
|
||||
|
||||
// Build the service with a SecretBoxService stub whose decryptSecret is supplied
|
||||
// per-test. The repo dep is unused by the methods under test.
|
||||
function buildService(decryptSecret: (blob: string) => string) {
|
||||
const secretBox = { decryptSecret: jest.fn(decryptSecret) };
|
||||
const service = new McpClientsService({} as never, secretBox as never);
|
||||
return { service, secretBox };
|
||||
}
|
||||
|
||||
describe('McpClientsService.decryptHeaders', () => {
|
||||
// Reach the private method via the as-any pattern common in these NestJS specs.
|
||||
const callDecrypt = (
|
||||
service: McpClientsService,
|
||||
blob: string | null,
|
||||
): Record<string, string> | undefined =>
|
||||
(
|
||||
service as unknown as {
|
||||
decryptHeaders: (b: string | null) => Record<string, string> | undefined;
|
||||
}
|
||||
).decryptHeaders(blob);
|
||||
|
||||
it('returns undefined for a null blob without decrypting', () => {
|
||||
const { service, secretBox } = buildService(() => '{}');
|
||||
expect(callDecrypt(service, null)).toBeUndefined();
|
||||
expect(secretBox.decryptSecret).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('decrypts a valid blob and keeps only string-valued headers', () => {
|
||||
const { service } = buildService(() =>
|
||||
JSON.stringify({
|
||||
Authorization: 'Bearer abc',
|
||||
'X-Api-Key': 'k',
|
||||
// Non-string values must be dropped, not coerced.
|
||||
count: 5,
|
||||
flag: true,
|
||||
nested: { a: 1 },
|
||||
}),
|
||||
);
|
||||
expect(callDecrypt(service, 'cipher')).toEqual({
|
||||
Authorization: 'Bearer abc',
|
||||
'X-Api-Key': 'k',
|
||||
});
|
||||
});
|
||||
|
||||
it('returns undefined when the decrypted object has no string headers', () => {
|
||||
const { service } = buildService(() => JSON.stringify({ count: 5 }));
|
||||
// No usable headers -> undefined (connect with no auth header), not {}.
|
||||
expect(callDecrypt(service, 'cipher')).toBeUndefined();
|
||||
});
|
||||
|
||||
it('FAILS OPEN: a decrypt error returns undefined instead of throwing', () => {
|
||||
const { service } = buildService(() => {
|
||||
throw new Error('Failed to decrypt secret — APP_SECRET may have changed');
|
||||
});
|
||||
const warnSpy = jest
|
||||
.spyOn(
|
||||
(service as unknown as { logger: { warn: (...a: unknown[]) => void } })
|
||||
.logger,
|
||||
'warn',
|
||||
)
|
||||
.mockImplementation(() => undefined);
|
||||
|
||||
let result: unknown;
|
||||
expect(() => {
|
||||
result = callDecrypt(service, 'tampered-blob');
|
||||
}).not.toThrow();
|
||||
expect(result).toBeUndefined();
|
||||
// It warns (so ops sees degradation) but never logs the blob itself.
|
||||
expect(warnSpy).toHaveBeenCalledTimes(1);
|
||||
expect(String(warnSpy.mock.calls[0]?.[0])).not.toContain('tampered-blob');
|
||||
});
|
||||
|
||||
it('FAILS OPEN: malformed JSON (decrypts to non-JSON) returns undefined', () => {
|
||||
const { service } = buildService(() => 'not-json{');
|
||||
jest
|
||||
.spyOn(
|
||||
(service as unknown as { logger: { warn: (...a: unknown[]) => void } })
|
||||
.logger,
|
||||
'warn',
|
||||
)
|
||||
.mockImplementation(() => undefined);
|
||||
expect(callDecrypt(service, 'cipher')).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('McpClientsService.guardedFetch (SSRF per-request guard)', () => {
|
||||
// The bound guardedFetch closure lives on the instance as a private field.
|
||||
const guardedFetchOf = (service: McpClientsService) =>
|
||||
(service as unknown as { guardedFetch: typeof fetch }).guardedFetch;
|
||||
|
||||
let fetchSpy: jest.SpiedFunction<typeof fetch>;
|
||||
|
||||
beforeEach(() => {
|
||||
// Any reachable real fetch would be a network call; assert per-test that the
|
||||
// blocked paths never reach it, and stub a Response for the allow path.
|
||||
fetchSpy = jest
|
||||
.spyOn(global, 'fetch')
|
||||
.mockResolvedValue(new Response('ok', { status: 200 }));
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
jest.restoreAllMocks();
|
||||
});
|
||||
|
||||
const blocked: Array<[string, string]> = [
|
||||
['loopback IPv4', 'http://127.0.0.1/mcp'],
|
||||
['private 10/8', 'http://10.0.0.5/mcp'],
|
||||
['private 192.168/16', 'http://192.168.1.1/mcp'],
|
||||
['cloud metadata link-local', 'http://169.254.169.254/latest/meta-data/'],
|
||||
['loopback IPv6 (bracketed)', 'http://[::1]:8080/mcp'],
|
||||
];
|
||||
|
||||
it.each(blocked)(
|
||||
'rejects a request to %s without opening a socket',
|
||||
async (_label, url) => {
|
||||
const { service } = buildService(() => '{}');
|
||||
await expect(guardedFetchOf(service)(url)).rejects.toThrow(
|
||||
/blocked request/,
|
||||
);
|
||||
expect(fetchSpy).not.toHaveBeenCalled();
|
||||
},
|
||||
);
|
||||
|
||||
it('rejects an unparseable URL as a blocked request', async () => {
|
||||
const { service } = buildService(() => '{}');
|
||||
await expect(
|
||||
guardedFetchOf(service)('::: not a url :::'),
|
||||
).rejects.toThrow('blocked request: invalid URL');
|
||||
expect(fetchSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('allows a public IP literal and forwards through the pinned dispatcher', async () => {
|
||||
const { service } = buildService(() => '{}');
|
||||
const res = await guardedFetchOf(service)('http://8.8.8.8/mcp');
|
||||
|
||||
expect(res.status).toBe(200);
|
||||
expect(fetchSpy).toHaveBeenCalledTimes(1);
|
||||
// The init MUST carry the SSRF-pinned undici dispatcher (the rebinding pin);
|
||||
// dropping it would let undici do a second, unchecked DNS resolution.
|
||||
const init = fetchSpy.mock.calls[0][1] as RequestInit & {
|
||||
dispatcher?: unknown;
|
||||
};
|
||||
expect(init.dispatcher).toBeDefined();
|
||||
});
|
||||
});
|
||||
@@ -6,6 +6,34 @@ import { esmImport } from '../../../common/helpers/esm-import';
|
||||
* ESM-only `@docmost/mcp` package. We only need the constructor + the read/write
|
||||
* methods used by the per-user tool adapter; the full client surface lives in
|
||||
* `packages/mcp/src/client.ts`. Signatures here mirror that file exactly.
|
||||
*
|
||||
* DRIFT GUARD: the method NAMES below are runtime-checked against the real
|
||||
* `DocmostClient` by `packages/mcp/test/unit/client-host-contract.test.mjs`
|
||||
* (which can import the ESM class directly). If you rename/remove a method here
|
||||
* or in client.ts, that test fails — so a stale mirror cannot silently ship a
|
||||
* runtime "x is not a function" into an agent tool call. Keep the two in sync.
|
||||
*
|
||||
* STAGED PLAN — full derivation `DocmostClientLike = <real DocmostClient type>`
|
||||
* (issue #193, layer 3) is intentionally NOT done; it stays a hand-mirror for
|
||||
* now because of two verified blockers across the ESM(mcp)/CJS(server) boundary:
|
||||
* 1. `@docmost/mcp` emits NO declaration files (its tsconfig has no
|
||||
* `declaration`, package.json has no `types`/types-export) and the server
|
||||
* tsconfig has no path mapping for it — the server only loads it via the
|
||||
* runtime `import()` trick below, so there is no type to import today.
|
||||
* 2. The real client methods have inferred, CONCRETE return types; the in-app
|
||||
* tool adapter reads results through loose `Record<string,unknown>` returns
|
||||
* + `as` casts (e.g. `(result?.data ?? {}) as { title?: string }`).
|
||||
* Deriving the exact type would make those casts non-overlapping ("may be a
|
||||
* mistake") and break the build, and `Partial<DocmostClientLike>` test stubs
|
||||
* would have to satisfy the full concrete surface.
|
||||
* To do it safely later (incrementally): (a) turn on `declaration: true` in
|
||||
* packages/mcp/tsconfig.json + add a `types` export condition and commit the
|
||||
* emitted `.d.ts`; (b) `import type { DocmostClient } from '@docmost/mcp'` here
|
||||
* and replace this interface with a `Pick<DocmostClient, ...>` of the consumed
|
||||
* methods; (c) audit every `as` cast in ai-chat-tools.service.ts against the now
|
||||
* concrete return types (double-cast through `unknown` only where genuinely
|
||||
* needed); (d) keep the runtime guard test as a belt-and-braces check. Until
|
||||
* then the guard test above is the cheap, behaviour-neutral protection.
|
||||
*/
|
||||
export interface DocmostClientLike {
|
||||
// --- read ---
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
import { z } from 'zod';
|
||||
import { AiChatToolsService } from './ai-chat-tools.service';
|
||||
import * as loader from './docmost-client.loader';
|
||||
import type { DocmostClientLike } from './docmost-client.loader';
|
||||
// The real zod-agnostic registry, imported from source so the contract is checked
|
||||
// against exactly what the @docmost/mcp package ships (no hand-stub).
|
||||
import { SHARED_TOOL_SPECS } from '../../../../../../packages/mcp/src/tool-specs';
|
||||
|
||||
/**
|
||||
* CONTRACT: SHARED_TOOL_SPECS <-> in-app tool wiring parity.
|
||||
*
|
||||
* `packages/mcp/src/tool-specs.ts` is the single source of truth for the tools
|
||||
* that are intentionally IDENTICAL across the standalone MCP server (zod v3) and
|
||||
* the in-app AI-SDK service (zod v4). The in-app service builds each one via
|
||||
* `sharedTool(sharedToolSpecs.<key>, execute)`, keyed by the spec's `inAppKey`.
|
||||
*
|
||||
* This test fails the build if a spec is added to the registry but never wired
|
||||
* in-app, if an `inAppKey` is renamed without updating the service, if the
|
||||
* description drifts between the registry and the exposed tool, if the
|
||||
* snake_case `mcpName` <-> camelCase `inAppKey` convention is broken, or if the
|
||||
* exposed tool's input-schema keys diverge from the spec's `buildShape`.
|
||||
*
|
||||
* It does NOT need @docmost/mcp built: the registry is imported from TS source,
|
||||
* and the ESM loader is mocked so `forUser()` never dynamically imports the
|
||||
* package.
|
||||
*/
|
||||
describe('SHARED_TOOL_SPECS contract parity', () => {
|
||||
// Empty fake client: no tool is executed here — every assertion is on tool
|
||||
// presence / metadata / schema, so the client methods are never called.
|
||||
const fakeClient: Partial<DocmostClientLike> = {};
|
||||
const tokenServiceStub = {
|
||||
generateAccessToken: jest.fn().mockResolvedValue('access-token'),
|
||||
generateCollabToken: jest.fn().mockResolvedValue('collab-token'),
|
||||
};
|
||||
|
||||
let tools: Record<string, unknown>;
|
||||
|
||||
beforeAll(async () => {
|
||||
jest.spyOn(loader, 'loadDocmostMcp').mockResolvedValue({
|
||||
DocmostClient: function () {
|
||||
return fakeClient as DocmostClientLike;
|
||||
} as unknown as loader.DocmostClientCtor,
|
||||
// Feed the service the SAME registry this test asserts against.
|
||||
sharedToolSpecs: SHARED_TOOL_SPECS as unknown as Record<
|
||||
string,
|
||||
loader.SharedToolSpec
|
||||
>,
|
||||
});
|
||||
const 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,
|
||||
);
|
||||
tools = (await service.forUser(
|
||||
{ id: 'user-1', email: 'u@example.com', workspaceId: 'ws-1' } as never,
|
||||
'session-1',
|
||||
'ws-1',
|
||||
'chat-1',
|
||||
)) as unknown as Record<string, unknown>;
|
||||
});
|
||||
|
||||
afterAll(() => jest.restoreAllMocks());
|
||||
|
||||
// camelCase -> snake_case, matching the registry's mcpName convention.
|
||||
const toSnake = (s: string) =>
|
||||
s.replace(/[A-Z]/g, (c) => `_${c.toLowerCase()}`);
|
||||
|
||||
// Type as the (optional-buildShape) SharedToolSpec; the `satisfies` literal
|
||||
// above otherwise narrows to a union where some members lack buildShape.
|
||||
const specEntries = Object.entries(SHARED_TOOL_SPECS) as Array<
|
||||
[string, loader.SharedToolSpec]
|
||||
>;
|
||||
|
||||
// Sanity: the registry is non-empty, so the per-spec table below is not vacuous.
|
||||
it('registry is non-empty', () => {
|
||||
expect(specEntries.length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
describe.each(specEntries)('spec "%s"', (registryKey, spec) => {
|
||||
it('registry key equals its inAppKey', () => {
|
||||
// The service indexes the registry by property name; a key != inAppKey
|
||||
// would wire the wrong (or no) tool.
|
||||
expect(spec.inAppKey).toBe(registryKey);
|
||||
});
|
||||
|
||||
it('mcpName is the snake_case form of inAppKey', () => {
|
||||
expect(spec.mcpName).toBe(toSnake(spec.inAppKey));
|
||||
});
|
||||
|
||||
it('is exposed in-app under its inAppKey', () => {
|
||||
// Fails if a spec is added to the registry but never wired in forUser().
|
||||
expect(tools[spec.inAppKey]).toBeDefined();
|
||||
});
|
||||
|
||||
it("exposed tool's description matches the registry description", () => {
|
||||
const tool = tools[spec.inAppKey] as { description: string };
|
||||
expect(tool.description).toBe(spec.description);
|
||||
});
|
||||
|
||||
it("exposed tool's input-schema keys match buildShape (incl. required)", () => {
|
||||
const tool = tools[spec.inAppKey] as {
|
||||
inputSchema: { jsonSchema: { properties?: Record<string, unknown>; required?: string[] } };
|
||||
};
|
||||
const json = tool.inputSchema.jsonSchema;
|
||||
const actualKeys = Object.keys(json.properties ?? {}).sort();
|
||||
|
||||
// Derive the spec's declared shape with THIS layer's zod (v4) — the same
|
||||
// call the service makes — then compare key sets and required-ness.
|
||||
const shape = spec.buildShape ? spec.buildShape(z) : {};
|
||||
const expectedKeys = Object.keys(shape).sort();
|
||||
expect(actualKeys).toEqual(expectedKeys);
|
||||
|
||||
// A non-.optional() field must surface as required in the advertised schema.
|
||||
const expectedRequired = Object.entries(shape)
|
||||
.filter(([, field]) => !(field as z.ZodTypeAny).isOptional?.())
|
||||
.map(([k]) => k)
|
||||
.sort();
|
||||
expect((json.required ?? []).slice().sort()).toEqual(expectedRequired);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,129 @@
|
||||
import { ShareService } from './share.service';
|
||||
|
||||
// Sibling of share-comment-strip.spec.ts. The public-share sanitizer strips ONLY
|
||||
// `comment` marks (internal-team metadata) via removeMarkTypeFromDoc(doc,
|
||||
// 'comment'). The `spoiler` mark is legitimate authored content (hidden text the
|
||||
// reader clicks to reveal) and MUST survive the share-strip — otherwise public
|
||||
// readers would see the secret in plain text or lose it entirely.
|
||||
//
|
||||
// We drive the SAME real seam the comment-strip test uses:
|
||||
// updatePublicAttachments -> prepareContentForShare -> removeMarkTypeFromDoc.
|
||||
|
||||
const WS = 'ws-1';
|
||||
const PAGE = 'page-1';
|
||||
|
||||
function buildService() {
|
||||
const shareRepo = { findById: jest.fn() };
|
||||
const pageRepo = { findById: jest.fn() };
|
||||
const pagePermissionRepo = {
|
||||
hasRestrictedAncestor: jest.fn(async () => false),
|
||||
};
|
||||
const tokenService = {
|
||||
generateAttachmentToken: jest.fn(async () => 'tok'),
|
||||
};
|
||||
const workspaceRepo = {
|
||||
findById: jest.fn(async () => ({ id: WS, settings: { htmlEmbed: true } })),
|
||||
};
|
||||
|
||||
return new ShareService(
|
||||
shareRepo as any,
|
||||
pageRepo as any,
|
||||
pagePermissionRepo as any,
|
||||
{} as any, // db (unused on this path)
|
||||
tokenService as any,
|
||||
{} as any, // transclusionService (unused)
|
||||
workspaceRepo as any,
|
||||
);
|
||||
}
|
||||
|
||||
// Text carrying a `spoiler` mark (no attributes; revealed state is UI-only).
|
||||
function spoilerText(text: string) {
|
||||
return {
|
||||
type: 'text',
|
||||
text,
|
||||
marks: [{ type: 'spoiler' }],
|
||||
};
|
||||
}
|
||||
|
||||
// Text carrying a `comment` mark with an id (the thing that DOES get stripped).
|
||||
function commentedText(text: string, commentId: string) {
|
||||
return {
|
||||
type: 'text',
|
||||
text,
|
||||
marks: [{ type: 'comment', attrs: { commentId, resolved: false } }],
|
||||
};
|
||||
}
|
||||
|
||||
async function sanitize(content: any) {
|
||||
const service = buildService();
|
||||
return service.updatePublicAttachments({
|
||||
id: PAGE,
|
||||
workspaceId: WS,
|
||||
content,
|
||||
} as any);
|
||||
}
|
||||
|
||||
function countMarks(doc: any, type: string): number {
|
||||
let count = 0;
|
||||
const walk = (node: any) => {
|
||||
if (!node || typeof node !== 'object') return;
|
||||
if (Array.isArray(node.marks)) {
|
||||
for (const mark of node.marks) {
|
||||
if (mark?.type === type) count++;
|
||||
}
|
||||
}
|
||||
if (Array.isArray(node.content)) node.content.forEach(walk);
|
||||
};
|
||||
walk(doc);
|
||||
return count;
|
||||
}
|
||||
|
||||
describe('ShareService keeps spoiler marks on public shares (real code)', () => {
|
||||
it('does NOT strip a spoiler mark', async () => {
|
||||
const content = {
|
||||
type: 'doc',
|
||||
content: [
|
||||
{
|
||||
type: 'paragraph',
|
||||
content: [{ type: 'text', text: 'visible ' }, spoilerText('hidden')],
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
expect(countMarks(content, 'spoiler')).toBe(1);
|
||||
|
||||
const out = await sanitize(content);
|
||||
|
||||
// The spoiler mark survives the share-strip.
|
||||
expect(countMarks(out, 'spoiler')).toBe(1);
|
||||
expect(JSON.stringify(out)).toContain('hidden');
|
||||
});
|
||||
|
||||
it('strips comment marks but keeps spoiler marks in the same doc', async () => {
|
||||
const content = {
|
||||
type: 'doc',
|
||||
content: [
|
||||
{
|
||||
type: 'paragraph',
|
||||
content: [
|
||||
commentedText('reviewed', 'cmt-1'),
|
||||
{ type: 'text', text: ' and ' },
|
||||
spoilerText('secret'),
|
||||
],
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
expect(countMarks(content, 'comment')).toBe(1);
|
||||
expect(countMarks(content, 'spoiler')).toBe(1);
|
||||
|
||||
const out = await sanitize(content);
|
||||
|
||||
// comment is removed, spoiler is preserved.
|
||||
expect(countMarks(out, 'comment')).toBe(0);
|
||||
expect(countMarks(out, 'spoiler')).toBe(1);
|
||||
const serialized = JSON.stringify(out);
|
||||
expect(serialized).not.toContain('cmt-1');
|
||||
expect(serialized).toContain('secret');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,167 @@
|
||||
import { PageRepo } from './page.repo';
|
||||
import {
|
||||
DummyDriver,
|
||||
Kysely,
|
||||
PostgresAdapter,
|
||||
PostgresIntrospector,
|
||||
PostgresQueryCompiler,
|
||||
} from 'kysely';
|
||||
|
||||
/**
|
||||
* F6 regression guard for the embeddable-page predicate.
|
||||
*
|
||||
* The predicate is shared by `countEmbeddablePages` (the "Indexed N of M" coverage
|
||||
* denominator) and `getEmbeddablePageIds` (the exact set a full reindex iterates).
|
||||
* It MUST select pages whose `text_content` was never backfilled (null/empty) but
|
||||
* whose ProseMirror `content` JSON still carries body text — `reindexPage` builds
|
||||
* its chunks straight from `content`, so without a content clause such a page is
|
||||
* silently SKIPPED by a mass reindex even though it is fully embeddable.
|
||||
*
|
||||
* The content clause keys on the structural text-node marker `"type":"text"`, NOT
|
||||
* a bare `"text":` key. The bare key also appears as the `attrs.text` of atom
|
||||
* nodes that carry NO extractable text — notably math (`mathBlock`/`mathInline`),
|
||||
* whose LaTeX lives in `attrs.text` and has no `generateText` serializer. A
|
||||
* math-ONLY page therefore yields empty `text_content` and zero embeddings; if the
|
||||
* predicate matched its `attrs.text` it would land in the denominator but
|
||||
* `reindexPage` would no-op on it, pinning "Indexed N of M" below 100% forever —
|
||||
* the exact bug this feature fixes. The `"type":"text"` marker matches only real
|
||||
* text nodes (what `jsonToText` extracts), keeping the predicate consistent with
|
||||
* what gets indexed.
|
||||
*
|
||||
* There is no real Postgres here: a recording Kysely (DummyDriver wired to the
|
||||
* Postgres query compiler) compiles the queries to SQL so we can assert the WHERE
|
||||
* predicate ORs in the narrowed content clause alongside the existing text_content
|
||||
* and stored-embeddings clauses — and that BOTH callers compile the identical
|
||||
* clause (denominator and reindex set can never diverge).
|
||||
*/
|
||||
function makeRecordingDb() {
|
||||
const sqls: string[] = [];
|
||||
const db = new Kysely<any>({
|
||||
dialect: {
|
||||
createAdapter: () => new PostgresAdapter(),
|
||||
createDriver: () =>
|
||||
new (class extends DummyDriver {
|
||||
async acquireConnection() {
|
||||
return {
|
||||
executeQuery: async (compiled: { sql: string }) => {
|
||||
sqls.push(compiled.sql);
|
||||
return { rows: [] };
|
||||
},
|
||||
// eslint-disable-next-line @typescript-eslint/no-empty-function
|
||||
streamQuery: async function* () {},
|
||||
} as any;
|
||||
}
|
||||
})(),
|
||||
createIntrospector: (d: Kysely<any>) => new PostgresIntrospector(d),
|
||||
createQueryCompiler: () => new PostgresQueryCompiler(),
|
||||
},
|
||||
});
|
||||
return { db, sqls };
|
||||
}
|
||||
|
||||
// The narrowed content clause, as it appears in the compiled SQL. Keying on the
|
||||
// structural `"type":"text"` marker (not a bare `"text":` key) is what excludes
|
||||
// math-only pages whose only `"text"` key is the atom node's `attrs.text`.
|
||||
const NARROWED_CLAUSE = `"type"[[:space:]]*:[[:space:]]*"text"`;
|
||||
const BARE_TEXT_KEY = `"text"[[:space:]]*:`;
|
||||
|
||||
describe('PageRepo embeddable predicate — content-bearing pages (F6)', () => {
|
||||
it('selects content-bearing pages via the narrowed "type":"text" node marker', async () => {
|
||||
const { db, sqls } = makeRecordingDb();
|
||||
const repo = new PageRepo(db as any, {} as any, { emit: jest.fn() } as any);
|
||||
|
||||
await repo.getEmbeddablePageIds('ws-1');
|
||||
|
||||
expect(sqls).toHaveLength(1);
|
||||
const sql = sqls[0];
|
||||
|
||||
// Clause 1 (existing): pages with extractable text_content.
|
||||
expect(sql).toContain('text_content');
|
||||
// Clause 3 (the F6 fix, now narrowed): a page whose content JSON carries a
|
||||
// real text node is selected even when text_content is null/empty, so a full
|
||||
// reindex visits it instead of silently skipping it.
|
||||
expect(sql).toContain('content::text');
|
||||
expect(sql).toContain(NARROWED_CLAUSE);
|
||||
// It must NOT use the old bare `"text":` key, which also matches the
|
||||
// `attrs.text` of math-only atom pages (false-positive denominator inflation).
|
||||
expect(sql).not.toContain(BARE_TEXT_KEY);
|
||||
// Clause 2 (existing): pages that already have stored embeddings stay in the
|
||||
// set so a reindex can clear their stale rows.
|
||||
expect(sql.toLowerCase()).toContain('embeddings');
|
||||
});
|
||||
|
||||
it('countEmbeddablePages compiles the SAME narrowed clause as getEmbeddablePageIds', async () => {
|
||||
// Consistency is the core requirement: the denominator (countEmbeddablePages)
|
||||
// and the reindex set (getEmbeddablePageIds) MUST share the identical
|
||||
// predicate, else the live "done" counter and the steady-state total diverge.
|
||||
const { db, sqls } = makeRecordingDb();
|
||||
const repo = new PageRepo(db as any, {} as any, { emit: jest.fn() } as any);
|
||||
|
||||
await repo.countEmbeddablePages('ws-1');
|
||||
await repo.getEmbeddablePageIds('ws-1');
|
||||
|
||||
expect(sqls).toHaveLength(2);
|
||||
const [countSql, idsSql] = sqls;
|
||||
|
||||
// Both carry the narrowed content clause...
|
||||
expect(countSql).toContain(NARROWED_CLAUSE);
|
||||
expect(idsSql).toContain(NARROWED_CLAUSE);
|
||||
// ...neither carries the bare key...
|
||||
expect(countSql).not.toContain(BARE_TEXT_KEY);
|
||||
expect(idsSql).not.toContain(BARE_TEXT_KEY);
|
||||
// ...and the full OR predicate (text_content + content node + embeddings
|
||||
// EXISTS) is byte-identical between the two queries, so they can't drift.
|
||||
const where = (s: string) => s.slice(s.indexOf('where'));
|
||||
expect(where(countSql)).toEqual(where(idsSql));
|
||||
});
|
||||
|
||||
it('the content regex matches a text-bearing doc but NOT a math-only doc', () => {
|
||||
// Semantic check of the predicate against sample `content::text` payloads.
|
||||
// Note: `jsonb::text` is NOT identical to JSON.stringify — Postgres renders a
|
||||
// space after each colon (`"type": "text"`), which is exactly why the POSIX
|
||||
// clause uses `[[:space:]]*`. The clause `"type"[[:space:]]*:[[:space:]]*"text"`
|
||||
// maps to the JS regex below (`[[:space:]]` -> `\s`, tolerating both forms);
|
||||
// we evaluate it the way Postgres would.
|
||||
const re = /"type"\s*:\s*"text"/;
|
||||
|
||||
// A real paragraph with a text node -> embeddable.
|
||||
const textDoc = JSON.stringify({
|
||||
type: 'doc',
|
||||
content: [
|
||||
{
|
||||
type: 'paragraph',
|
||||
content: [{ type: 'text', text: 'hello world' }],
|
||||
},
|
||||
],
|
||||
});
|
||||
// A doc whose ONLY node is a math atom. Its LaTeX is in `attrs.text`, there is
|
||||
// no text node, and `jsonToText`/`generateText` has no serializer for it -> it
|
||||
// yields empty text_content and zero embeddings, so it must NOT qualify.
|
||||
const mathOnlyDoc = JSON.stringify({
|
||||
type: 'doc',
|
||||
content: [
|
||||
{ type: 'mathBlock', attrs: { text: 'E = mc^2' } },
|
||||
{ type: 'mathInline', attrs: { text: '\\alpha' } },
|
||||
],
|
||||
});
|
||||
// An empty doc has no text node either.
|
||||
const emptyDoc = JSON.stringify({ type: 'doc', content: [] });
|
||||
|
||||
expect(re.test(textDoc)).toBe(true);
|
||||
expect(re.test(mathOnlyDoc)).toBe(false);
|
||||
expect(re.test(emptyDoc)).toBe(false);
|
||||
// Sanity: the OLD bare-key regex WOULD have wrongly matched the math-only doc,
|
||||
// which is precisely the false positive the narrowing removes.
|
||||
expect(/"text"\s*:/.test(mathOnlyDoc)).toBe(true);
|
||||
|
||||
// A user literally TYPING `"type":"text"` in prose can't false-positive on an
|
||||
// otherwise text-less page: in `content::text` the typed value's quotes are
|
||||
// escaped (`\"type\":\"text\"`), so the literal-quote regex does not match the
|
||||
// escaped form. (And such a page is a genuine text node anyway.)
|
||||
const escapedLiteral = JSON.stringify({
|
||||
type: 'doc',
|
||||
content: [{ type: 'someAtom', attrs: { note: '"type":"text"' } }],
|
||||
});
|
||||
expect(re.test(escapedLiteral)).toBe(false);
|
||||
});
|
||||
});
|
||||
@@ -12,6 +12,7 @@ import { executeWithCursorPagination } from '@docmost/db/pagination/cursor-pagin
|
||||
import { validate as isValidUUID } from 'uuid';
|
||||
import { ExpressionBuilder, sql } from 'kysely';
|
||||
import { DB } from '@docmost/db/types/db';
|
||||
import { DbInterface } from '@docmost/db/types/db.interface';
|
||||
import { jsonArrayFrom, jsonObjectFrom } from 'kysely/helpers/postgres';
|
||||
import { SpaceMemberRepo } from '@docmost/db/repos/space/space-member.repo';
|
||||
import { EventEmitter2 } from '@nestjs/event-emitter';
|
||||
@@ -233,9 +234,9 @@ export class PageRepo {
|
||||
* text-less pages (which legitimately store zero embeddings) don't keep the
|
||||
* bar below 100% forever.
|
||||
*
|
||||
* A page qualifies if it has non-empty textContent OR already has stored
|
||||
* embeddings. The second clause covers pages whose text the indexer extracted
|
||||
* from the content JSON when textContent was null, and guarantees this total is
|
||||
* A page qualifies if it has non-empty textContent, OR its content JSON has at
|
||||
* least one text node (`"type":"text"`) when textContent was never backfilled,
|
||||
* OR it already has stored embeddings. The last clause guarantees this total is
|
||||
* always >= countIndexedPages (the indexed count can never exceed it).
|
||||
*/
|
||||
async countEmbeddablePages(workspaceId: string): Promise<number> {
|
||||
@@ -243,37 +244,91 @@ export class PageRepo {
|
||||
.selectFrom('pages as p')
|
||||
.where('p.workspaceId', '=', workspaceId)
|
||||
.where('p.deletedAt', 'is', null)
|
||||
.where((eb) =>
|
||||
eb.or([
|
||||
// Has extractable body text. The regex matches any non-whitespace
|
||||
// character, mirroring the indexer's `text.trim().length === 0` check
|
||||
// (raw SQL -> use the snake_case column name).
|
||||
sql<boolean>`p.text_content ~ '[^[:space:]]'`,
|
||||
// OR already has at least one (non-deleted) embedding row.
|
||||
eb.exists(
|
||||
eb
|
||||
.selectFrom('pageEmbeddings as pe')
|
||||
.select(sql`1`.as('one'))
|
||||
.whereRef('pe.pageId', '=', 'p.id')
|
||||
.where('pe.deletedAt', 'is', null),
|
||||
),
|
||||
]),
|
||||
)
|
||||
.where((eb) => this.embeddablePredicate(eb))
|
||||
.select((eb) => eb.fn.countAll().as('count'))
|
||||
.executeTakeFirst();
|
||||
return Number(row?.count ?? 0);
|
||||
}
|
||||
|
||||
/**
|
||||
* IDs of all non-deleted pages in a workspace. Used by the RAG bulk reindex to
|
||||
* (re)build embeddings for every existing page.
|
||||
* The "embeddable content" qualifying predicate, shared verbatim by
|
||||
* countEmbeddablePages (the steady-state denominator) and getEmbeddablePageIds
|
||||
* (the set the bulk reindex iterates). Both MUST use the exact same condition
|
||||
* or the live total and steady-state total diverge — extracting it here is what
|
||||
* guarantees that, replacing the previous hand-duplicated copy. Callers supply
|
||||
* the trivial workspaceId/deletedAt filters inline; this returns only the
|
||||
* non-trivial OR clause, evaluated against the `p` alias of `pages`.
|
||||
*
|
||||
* A page qualifies if it has non-empty textContent, OR its ProseMirror
|
||||
* `content` JSON has at least one text node (`"type":"text"`) even though
|
||||
* textContent was never backfilled, OR it already has a stored (non-deleted)
|
||||
* embedding row.
|
||||
*/
|
||||
async getIdsByWorkspace(workspaceId: string): Promise<string[]> {
|
||||
private embeddablePredicate(
|
||||
eb: ExpressionBuilder<DbInterface & { p: DbInterface['pages'] }, 'p'>,
|
||||
) {
|
||||
return eb.or([
|
||||
// Has extractable body text. The regex matches any non-whitespace
|
||||
// character, mirroring the indexer's `text.trim().length === 0` check
|
||||
// (raw SQL -> use the snake_case column name).
|
||||
sql<boolean>`p.text_content ~ '[^[:space:]]'`,
|
||||
// OR the ProseMirror `content` JSON has at least one text node (`"type":
|
||||
// "text"`) the indexer can extract, even when `text_content` is null/empty
|
||||
// (never backfilled): `reindexPage` runs `jsonToText` (generateText) over
|
||||
// `content`, which only emits the text of ProseMirror text nodes, so such a
|
||||
// page IS embeddable and a full reindex MUST visit it (otherwise it is
|
||||
// silently skipped). A text node always serialises as
|
||||
// `{"type":"text","text":"..."}`, so we key on the structural `"type":
|
||||
// "text"` marker — NOT a bare `"text":` key, which also appears as the
|
||||
// `attrs.text` of atom nodes that carry NO extractable text (e.g. math
|
||||
// `mathBlock`/`mathInline`, whose LaTeX lives in `attrs.text` and has no
|
||||
// text serializer). A math-only page thus produces empty `text_content` and
|
||||
// zero embeddings; matching its `attrs.text` here would wrongly inflate the
|
||||
// denominator and keep "Indexed N of M" below 100% forever. An empty doc
|
||||
// (no text nodes) has no `"type":"text"` and is correctly excluded. A user
|
||||
// who literally types `"type":"text"` in their prose can't false-positive:
|
||||
// in `content::text` that text value's quotes are escaped (`\"type\"...`),
|
||||
// so the literal-quote regex won't match the escaped form (and such a page
|
||||
// is a real text node anyway).
|
||||
sql<boolean>`p.content::text ~ '"type"[[:space:]]*:[[:space:]]*"text"'`,
|
||||
// OR already has at least one (non-deleted) embedding row.
|
||||
eb.exists(
|
||||
eb
|
||||
.selectFrom('pageEmbeddings as pe')
|
||||
.select(sql`1`.as('one'))
|
||||
.whereRef('pe.pageId', '=', 'p.id')
|
||||
.where('pe.deletedAt', 'is', null),
|
||||
),
|
||||
]);
|
||||
}
|
||||
|
||||
/**
|
||||
* IDs of the EMBEDDABLE page set for a workspace — the exact same set that
|
||||
* `countEmbeddablePages` counts (a page qualifies if it has non-empty
|
||||
* textContent, OR content JSON with at least one text node (`"type":"text"`)
|
||||
* and an empty/null textContent, OR already has a stored embedding row). The
|
||||
* bulk reindex
|
||||
* iterates THIS set so the live "done" counter reaches exactly
|
||||
* `countEmbeddablePages` (the steady-state denominator), instead of iterating
|
||||
* every non-deleted page (which would push the denominator above the
|
||||
* steady-state value mid-run).
|
||||
*
|
||||
* IMPORTANT: the qualifying WHERE is shared with `countEmbeddablePages` via the
|
||||
* private `embeddablePredicate` helper, so the two can no longer drift — if the
|
||||
* embeddable definition changes, change it once there and both stay in lockstep
|
||||
* (else the live total and steady-state total diverge again). Dropping
|
||||
* text-less pages is correct: `reindexPage` no-ops on
|
||||
* a page with no extractable content anyway, and a page that lost its text but
|
||||
* still has stale embeddings IS in this set (the EXISTS clause), so it is still
|
||||
* visited and its stale rows are cleared.
|
||||
*/
|
||||
async getEmbeddablePageIds(workspaceId: string): Promise<string[]> {
|
||||
const rows = await this.db
|
||||
.selectFrom('pages')
|
||||
.select('id')
|
||||
.where('workspaceId', '=', workspaceId)
|
||||
.where('deletedAt', 'is', null)
|
||||
.selectFrom('pages as p')
|
||||
.select('p.id')
|
||||
.where('p.workspaceId', '=', workspaceId)
|
||||
.where('p.deletedAt', 'is', null)
|
||||
.where((eb) => this.embeddablePredicate(eb))
|
||||
.execute();
|
||||
return rows.map((r) => r.id);
|
||||
}
|
||||
|
||||
@@ -1,4 +1,12 @@
|
||||
import { parsePositiveInt } from './ai-settings.service';
|
||||
import { AiSettingsService, parsePositiveInt } from './ai-settings.service';
|
||||
import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';
|
||||
import { AiAgentRoleRepo } from '@docmost/db/repos/ai-agent-roles/ai-agent-roles.repo';
|
||||
import { AiProviderCredentialsRepo } from '@docmost/db/repos/ai-chat/ai-provider-credentials.repo';
|
||||
import { PageEmbeddingRepo } from '@docmost/db/repos/ai-chat/page-embedding.repo';
|
||||
import { PageRepo } from '@docmost/db/repos/page/page.repo';
|
||||
import { SecretBoxService } from '../crypto/secret-box';
|
||||
import { EmbeddingReindexProgressService } from './embedding-reindex-progress.service';
|
||||
import type { Queue } from 'bullmq';
|
||||
|
||||
/**
|
||||
* Round-trip coercion for numeric `::text` provider settings (e.g.
|
||||
@@ -41,3 +49,196 @@ describe('parsePositiveInt', () => {
|
||||
expect(parsePositiveInt(42)).toBe(42);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* getMasked must surface the LIVE reindex run progress while a reindex is active
|
||||
* (so the "Indexed X of Y" counter can climb 0 -> total), and fall back to the
|
||||
* steady-state DB coverage count (countIndexedPages / countEmbeddablePages) when
|
||||
* no reindex is running. This is the server side of the fix for the counter that
|
||||
* otherwise stays stuck at "478 of 478" the whole reindex.
|
||||
*/
|
||||
describe('AiSettingsService.getMasked reindex progress', () => {
|
||||
const WORKSPACE_ID = 'ws-1';
|
||||
|
||||
function makeService() {
|
||||
// No driver configured -> the credentials lookup is skipped, keeping the
|
||||
// setup minimal; we only care about the indexed/total numbers here.
|
||||
const workspaceRepo = {
|
||||
findById: jest.fn().mockResolvedValue({ settings: {} }),
|
||||
};
|
||||
const aiAgentRoleRepo = {};
|
||||
const aiProviderCredentialsRepo = { find: jest.fn() };
|
||||
const pageEmbeddingRepo = {
|
||||
countIndexedPages: jest.fn().mockResolvedValue(478),
|
||||
};
|
||||
const pageRepo = {
|
||||
countEmbeddablePages: jest.fn().mockResolvedValue(478),
|
||||
};
|
||||
const secretBox = {};
|
||||
const reindexProgress = {
|
||||
get: jest.fn().mockResolvedValue(null),
|
||||
};
|
||||
const aiQueue = {};
|
||||
|
||||
const service = new AiSettingsService(
|
||||
workspaceRepo as unknown as WorkspaceRepo,
|
||||
aiAgentRoleRepo as unknown as AiAgentRoleRepo,
|
||||
aiProviderCredentialsRepo as unknown as AiProviderCredentialsRepo,
|
||||
pageEmbeddingRepo as unknown as PageEmbeddingRepo,
|
||||
pageRepo as unknown as PageRepo,
|
||||
secretBox as unknown as SecretBoxService,
|
||||
reindexProgress as unknown as EmbeddingReindexProgressService,
|
||||
aiQueue as unknown as Queue,
|
||||
);
|
||||
return { service, reindexProgress, pageEmbeddingRepo };
|
||||
}
|
||||
|
||||
it('reports the live run numbers when a reindex progress record is active', async () => {
|
||||
const { service, reindexProgress } = makeService();
|
||||
// Use a progress.total (500) DISTINCT from the DB count (478) so the test
|
||||
// actually pins the progress.total branch rather than coincidentally
|
||||
// matching the DB fallback. With fix #1 the two sources agree in practice,
|
||||
// but getMasked must still return progress.total when a record is active.
|
||||
reindexProgress.get.mockResolvedValue({
|
||||
total: 500,
|
||||
done: 120,
|
||||
startedAt: Date.now(),
|
||||
});
|
||||
|
||||
const masked = await service.getMasked(WORKSPACE_ID);
|
||||
|
||||
expect(masked.indexedPages).toBe(120); // progress.done, not DB 478
|
||||
expect(masked.totalPages).toBe(500); // progress.total, not DB 478
|
||||
expect(masked.reindexing).toBe(true);
|
||||
});
|
||||
|
||||
it('falls back to countIndexedPages when no reindex is active', async () => {
|
||||
const { service, reindexProgress } = makeService();
|
||||
reindexProgress.get.mockResolvedValue(null);
|
||||
|
||||
const masked = await service.getMasked(WORKSPACE_ID);
|
||||
|
||||
expect(masked.indexedPages).toBe(478);
|
||||
expect(masked.totalPages).toBe(478);
|
||||
expect(masked.reindexing).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* reindex() must seed a live progress record (done=0) BEFORE enqueueing so the
|
||||
* first status poll shows 0 — but ONLY when no run is already active, since
|
||||
* aiQueue.add() de-duplicates a running reindex and a re-seed would reset the
|
||||
* visible counter to 0 while the live worker keeps incrementing from its real
|
||||
* position.
|
||||
*/
|
||||
describe('AiSettingsService.reindex progress seed', () => {
|
||||
const WORKSPACE_ID = 'ws-1';
|
||||
|
||||
function makeService() {
|
||||
const order: string[] = [];
|
||||
const aiQueue = {
|
||||
remove: jest.fn().mockResolvedValue(undefined),
|
||||
add: jest.fn().mockImplementation(async () => {
|
||||
order.push('add');
|
||||
}),
|
||||
};
|
||||
const pageRepo = {
|
||||
countEmbeddablePages: jest.fn().mockResolvedValue(478),
|
||||
};
|
||||
const reindexProgress = {
|
||||
// Default: no active run -> seed should happen.
|
||||
get: jest.fn().mockResolvedValue(null),
|
||||
start: jest.fn().mockImplementation(async () => {
|
||||
order.push('start');
|
||||
}),
|
||||
clear: jest.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
|
||||
const service = new AiSettingsService(
|
||||
{} as unknown as WorkspaceRepo,
|
||||
{} as unknown as AiAgentRoleRepo,
|
||||
{} as unknown as AiProviderCredentialsRepo,
|
||||
{} as unknown as PageEmbeddingRepo,
|
||||
pageRepo as unknown as PageRepo,
|
||||
{} as unknown as SecretBoxService,
|
||||
reindexProgress as unknown as EmbeddingReindexProgressService,
|
||||
aiQueue as unknown as Queue,
|
||||
);
|
||||
return { service, aiQueue, pageRepo, reindexProgress, order };
|
||||
}
|
||||
|
||||
it('seeds progress (workspace, count) BEFORE enqueue when no run is active', async () => {
|
||||
const { service, aiQueue, reindexProgress, order } = makeService();
|
||||
|
||||
await service.reindex(WORKSPACE_ID);
|
||||
|
||||
// The pre-seed carries the real page count AND a SHORT ttl (3rd arg) so a
|
||||
// de-duplicated enqueue against a just-finishing job can't leave a phantom
|
||||
// "reindexing: 0 of N" stuck for the full record TTL (F10).
|
||||
expect(reindexProgress.start).toHaveBeenCalledWith(
|
||||
WORKSPACE_ID,
|
||||
478,
|
||||
expect.any(Number),
|
||||
);
|
||||
const ttl = reindexProgress.start.mock.calls[0][2];
|
||||
// Short pre-seed TTL, distinct from the full 1h (3600s) record TTL, but
|
||||
// pinned to the client poll cap (120s) so a still-pending run can't expire
|
||||
// into a false "done" while the client is still polling (F11).
|
||||
expect(ttl).toBe(120);
|
||||
expect(aiQueue.add).toHaveBeenCalledTimes(1);
|
||||
// Seed must precede the enqueue so the first poll already reports done=0.
|
||||
expect(order).toEqual(['start', 'add']);
|
||||
});
|
||||
|
||||
it('does NOT re-seed when a run is already active (mid-run re-trigger)', async () => {
|
||||
const { service, aiQueue, reindexProgress } = makeService();
|
||||
// An active record exists -> a second click must not reset the counter.
|
||||
reindexProgress.get.mockResolvedValue({
|
||||
total: 478,
|
||||
done: 120,
|
||||
startedAt: Date.now(),
|
||||
});
|
||||
|
||||
await service.reindex(WORKSPACE_ID);
|
||||
|
||||
expect(reindexProgress.start).not.toHaveBeenCalled();
|
||||
// The enqueue still runs (and de-duplicates against the active job).
|
||||
expect(aiQueue.add).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('clears the seed it just wrote and re-throws when enqueue fails', async () => {
|
||||
const { service, aiQueue, reindexProgress } = makeService();
|
||||
// This call seeds (get() is null) but the enqueue then blows up
|
||||
// (Redis hiccup/shutdown) -> the worker never runs and never clear()s, so
|
||||
// reindex() must roll back its own seed to avoid a 1h stuck "reindexing".
|
||||
const boom = new Error('redis down');
|
||||
aiQueue.add.mockRejectedValue(boom);
|
||||
|
||||
await expect(service.reindex(WORKSPACE_ID)).rejects.toBe(boom);
|
||||
|
||||
expect(reindexProgress.start).toHaveBeenCalledWith(
|
||||
WORKSPACE_ID,
|
||||
478,
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(reindexProgress.clear).toHaveBeenCalledWith(WORKSPACE_ID);
|
||||
});
|
||||
|
||||
it('does NOT clear a concurrent active run when enqueue fails (no seed)', async () => {
|
||||
const { service, aiQueue, reindexProgress } = makeService();
|
||||
// A run is already active, so THIS call does not seed; if the enqueue then
|
||||
// fails it must NOT wipe the live worker's record.
|
||||
reindexProgress.get.mockResolvedValue({
|
||||
total: 478,
|
||||
done: 120,
|
||||
startedAt: Date.now(),
|
||||
});
|
||||
const boom = new Error('redis down');
|
||||
aiQueue.add.mockRejectedValue(boom);
|
||||
|
||||
await expect(service.reindex(WORKSPACE_ID)).rejects.toBe(boom);
|
||||
|
||||
expect(reindexProgress.start).not.toHaveBeenCalled();
|
||||
expect(reindexProgress.clear).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -8,6 +8,7 @@ import { AiProviderCredentialsRepo } from '@docmost/db/repos/ai-chat/ai-provider
|
||||
import { PageEmbeddingRepo } from '@docmost/db/repos/ai-chat/page-embedding.repo';
|
||||
import { PageRepo } from '@docmost/db/repos/page/page.repo';
|
||||
import { SecretBoxService } from '../crypto/secret-box';
|
||||
import { EmbeddingReindexProgressService } from './embedding-reindex-progress.service';
|
||||
import {
|
||||
AiDriver,
|
||||
AiProviderSettings,
|
||||
@@ -30,6 +31,30 @@ export function parsePositiveInt(raw: unknown): number | undefined {
|
||||
return Number.isFinite(n) && n > 0 ? Math.floor(n) : undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* TTL (seconds) for the enqueue-time progress PRE-SEED written by `reindex()`
|
||||
* before the worker starts. Deliberately SHORT relative to the full 1h record
|
||||
* TTL: if `aiQueue.add()` de-duplicates against a job that is just finishing
|
||||
* (the worker's finally already ran `clear()` but removeOnComplete hasn't yet
|
||||
* removed the job), no new worker runs to overwrite/clear this seed — so this
|
||||
* shorter TTL lets the phantom "reindexing: 0 of N" expire instead of sticking
|
||||
* for the full 1h record TTL. A worker that DOES start re-seeds with the full
|
||||
* TTL, so a real run is unaffected.
|
||||
*
|
||||
* It MUST be >= the client poll cap (REINDEX_POLL_CAP_MS = 120000ms in
|
||||
* ai-provider-settings.tsx) though: the AI_QUEUE worker runs at concurrency 1
|
||||
* and shares the queue with page-level embedding jobs, so a queued reindex can
|
||||
* wait well beyond a few dozen seconds before the worker re-seeds with the full
|
||||
* TTL. If the pre-seed expired while the job is still pending, `get()` returns
|
||||
* null and getMasked() falls back to the steady-state COUNT (indexedPages ==
|
||||
* totalPages, reindexing=false) — the client reads that as "done & fully
|
||||
* indexed", clears its deadline and STOPS polling, so the admin never sees the
|
||||
* real climb. Pinning the pre-seed TTL to the client cap means a deduped phantom
|
||||
* is bounded to ~120s — the same window the client already polls — and a genuine
|
||||
* pending run never expires-into-"done" inside that window.
|
||||
*/
|
||||
const PRE_SEED_TTL_SECONDS = 120;
|
||||
|
||||
/**
|
||||
* Shape of the partial update accepted by `update`. Mirrors the validated
|
||||
* controller DTO. `apiKey` / `embeddingApiKey` are write-only: undefined =
|
||||
@@ -74,6 +99,7 @@ export class AiSettingsService {
|
||||
private readonly pageEmbeddingRepo: PageEmbeddingRepo,
|
||||
private readonly pageRepo: PageRepo,
|
||||
private readonly secretBox: SecretBoxService,
|
||||
private readonly reindexProgress: EmbeddingReindexProgressService,
|
||||
@InjectQueue(QueueName.AI_QUEUE) private readonly aiQueue: Queue,
|
||||
) {}
|
||||
|
||||
@@ -100,21 +126,63 @@ export class AiSettingsService {
|
||||
.remove(`ai-search-disabled-${workspaceId}`)
|
||||
.catch(() => undefined);
|
||||
|
||||
// Seed a live progress record BEFORE enqueueing so the very first status
|
||||
// poll already reports done=0 (the reindex POST returns the PRE-job counts,
|
||||
// so without this seed the first poll would still show "total of total").
|
||||
// `totalPages` uses countEmbeddablePages — the SAME set the worker iterates
|
||||
// and the SAME denominator the status endpoint reports, so the live and
|
||||
// steady-state totals match.
|
||||
//
|
||||
// ONLY seed when no run is active: aiQueue.add() de-duplicates an already-
|
||||
// running reindex, so a mid-run re-trigger (second click / second admin /
|
||||
// second tab) must NOT reset the visible counter to 0 — that would
|
||||
// understate the live worker's real position for the rest of the run. The
|
||||
// worker's own start() at run begin is the single authoritative reset.
|
||||
let seeded = false;
|
||||
if ((await this.reindexProgress.get(workspaceId)) === null) {
|
||||
const totalPages = await this.pageRepo.countEmbeddablePages(workspaceId);
|
||||
// Short TTL (vs the full 1h record TTL): if add() below de-duplicates
|
||||
// against a just-finishing job whose worker already clear()ed but isn't
|
||||
// removed yet, no worker runs to clear this seed — the shorter TTL expires
|
||||
// the phantom record rather than leaving a stuck "reindexing: 0 of N" for
|
||||
// the full record TTL. It is kept >= the client poll cap (120s) so a
|
||||
// genuine but still-pending run never expires into a false "done" while
|
||||
// the client is still polling (see PRE_SEED_TTL_SECONDS).
|
||||
await this.reindexProgress.start(
|
||||
workspaceId,
|
||||
totalPages,
|
||||
PRE_SEED_TTL_SECONDS,
|
||||
);
|
||||
seeded = true;
|
||||
}
|
||||
|
||||
const jobId = `ai-reindex-${workspaceId}`;
|
||||
// Clear a prior non-active entry so a stale job can't block this reindex.
|
||||
// A locked/active job is left in place (remove() no-ops) and the add() below
|
||||
// de-duplicates against it, keeping the in-progress pass.
|
||||
await this.aiQueue.remove(jobId).catch(() => undefined);
|
||||
|
||||
await this.aiQueue.add(
|
||||
QueueJob.WORKSPACE_CREATE_EMBEDDINGS,
|
||||
{ workspaceId },
|
||||
{
|
||||
jobId,
|
||||
removeOnComplete: true,
|
||||
removeOnFail: true,
|
||||
},
|
||||
);
|
||||
try {
|
||||
await this.aiQueue.add(
|
||||
QueueJob.WORKSPACE_CREATE_EMBEDDINGS,
|
||||
{ workspaceId },
|
||||
{
|
||||
jobId,
|
||||
removeOnComplete: true,
|
||||
removeOnFail: true,
|
||||
},
|
||||
);
|
||||
} catch (err) {
|
||||
// If the enqueue fails (Redis hiccup/shutdown) the worker never runs, so
|
||||
// its finally->clear() never fires. Roll back the seed WE just wrote so
|
||||
// the status endpoint doesn't report a stuck "reindexing: 0 of N" for the
|
||||
// full TTL. Only clear when this call did the seed — never wipe a
|
||||
// concurrent active run's record (get() was non-null, seeded=false).
|
||||
if (seeded) {
|
||||
await this.reindexProgress.clear(workspaceId);
|
||||
}
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -253,13 +321,33 @@ export class AiSettingsService {
|
||||
hasSttApiKey = !!creds?.sttApiKeyEnc;
|
||||
}
|
||||
|
||||
// totalPages now counts only pages with embeddable content (non-empty text
|
||||
// or already-stored embeddings), so empty/text-less pages don't keep the
|
||||
// "Indexed N of M pages" bar below 100% forever.
|
||||
const [indexedPages, totalPages] = await Promise.all([
|
||||
this.pageEmbeddingRepo.countIndexedPages(workspaceId),
|
||||
this.pageRepo.countEmbeddablePages(workspaceId),
|
||||
]);
|
||||
// While a reindex run is active, report its LIVE progress (done climbs 0 ->
|
||||
// total) so the settings UI can watch it advance. Read progress FIRST and
|
||||
// short-circuit: this endpoint is polled every ~5s for the whole run, so when
|
||||
// a record is active we skip the two coverage COUNTs entirely (their results
|
||||
// would be discarded anyway). Without the live progress the counter never
|
||||
// drops: the per-page reindex hard-replaces rows in its own small
|
||||
// transaction, so countIndexedPages stays ~= total for the whole run. With no
|
||||
// active record we fall back to the steady-state DB coverage count, which
|
||||
// preserves the existing display and the client's "done == total -> stop
|
||||
// polling" condition (the run ends -> record cleared -> DB count == total).
|
||||
//
|
||||
// The fallback `totalPages` counts only pages with embeddable content
|
||||
// (non-empty text, content-borne text, or already-stored embeddings), so
|
||||
// empty/text-less pages don't keep the "Indexed N of M pages" bar below 100%
|
||||
// forever.
|
||||
const progress = await this.reindexProgress.get(workspaceId);
|
||||
let indexedPages: number;
|
||||
let totalPages: number;
|
||||
if (progress) {
|
||||
indexedPages = progress.done;
|
||||
totalPages = progress.total;
|
||||
} else {
|
||||
[indexedPages, totalPages] = await Promise.all([
|
||||
this.pageEmbeddingRepo.countIndexedPages(workspaceId),
|
||||
this.pageRepo.countEmbeddablePages(workspaceId),
|
||||
]);
|
||||
}
|
||||
|
||||
return {
|
||||
driver: provider.driver,
|
||||
@@ -281,6 +369,8 @@ export class AiSettingsService {
|
||||
hasSttApiKey,
|
||||
indexedPages,
|
||||
totalPages,
|
||||
// Optional hint for the client: a reindex run is currently in progress.
|
||||
reindexing: progress != null,
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ import { QueueName } from '../queue/constants';
|
||||
import { AiService } from './ai.service';
|
||||
import { AiSettingsService } from './ai-settings.service';
|
||||
import { AiSettingsController } from './ai-settings.controller';
|
||||
import { EmbeddingReindexProgressService } from './embedding-reindex-progress.service';
|
||||
|
||||
/**
|
||||
* LLM driver + provider-settings unit (§6.2/§6.4).
|
||||
@@ -19,7 +20,7 @@ import { AiSettingsController } from './ai-settings.controller';
|
||||
BullModule.registerQueue({ name: QueueName.AI_QUEUE }),
|
||||
],
|
||||
controllers: [AiSettingsController],
|
||||
providers: [AiService, AiSettingsService],
|
||||
exports: [AiService, AiSettingsService],
|
||||
providers: [AiService, AiSettingsService, EmbeddingReindexProgressService],
|
||||
exports: [AiService, AiSettingsService, EmbeddingReindexProgressService],
|
||||
})
|
||||
export class AiModule {}
|
||||
|
||||
@@ -146,4 +146,7 @@ export interface MaskedAiSettings {
|
||||
// RAG indexing coverage for the settings UI.
|
||||
indexedPages: number;
|
||||
totalPages: number;
|
||||
// True while a full workspace reindex is actively running (the counts above
|
||||
// then reflect the live run progress rather than the steady-state DB count).
|
||||
reindexing?: boolean;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
import { EmbeddingReindexProgressService } from './embedding-reindex-progress.service';
|
||||
import type { RedisService } from '@nestjs-labs/nestjs-ioredis';
|
||||
import type { Redis } from 'ioredis';
|
||||
|
||||
/**
|
||||
* Unit tests for the Redis-backed reindex-progress store.
|
||||
*
|
||||
* The store is a thin, BEST-EFFORT wrapper: writes (start/increment) issue an
|
||||
* hset/hincrby + expire pipeline and must SWALLOW Redis errors (progress is
|
||||
* cosmetic — it must never break a reindex); reads (get) must map a valid hash
|
||||
* to a ReindexProgress and degrade to null on a malformed/missing record or a
|
||||
* Redis failure. We drive it with a hand-rolled fake ioredis (the project mocks
|
||||
* Redis with plain fakes, see public-share limiter specs).
|
||||
*/
|
||||
describe('EmbeddingReindexProgressService', () => {
|
||||
const WORKSPACE_ID = 'ws-1';
|
||||
const KEY = 'ai:reindex:progress:ws-1';
|
||||
|
||||
/**
|
||||
* Build a fake ioredis whose `multi()` returns a chainable recorder and whose
|
||||
* `hgetall`/`del` are configurable jest mocks. `execImpl` lets a test make the
|
||||
* pipeline reject (to assert error-swallowing).
|
||||
*/
|
||||
function makeRedis(opts: { execImpl?: () => Promise<unknown> } = {}) {
|
||||
const exec = jest
|
||||
.fn()
|
||||
.mockImplementation(opts.execImpl ?? (() => Promise.resolve([])));
|
||||
// mockReturnThis() returns the call's `this` (the multi object), so the
|
||||
// chain hset().expire().exec() resolves correctly.
|
||||
const multiObj = {
|
||||
hset: jest.fn().mockReturnThis(),
|
||||
hincrby: jest.fn().mockReturnThis(),
|
||||
expire: jest.fn().mockReturnThis(),
|
||||
exec,
|
||||
};
|
||||
const multi = jest.fn(() => multiObj);
|
||||
const hgetall = jest.fn().mockResolvedValue({});
|
||||
const del = jest.fn().mockResolvedValue(1);
|
||||
const redis = { multi, hgetall, del } as unknown as Redis;
|
||||
return { redis, multiObj, multi, hgetall, del, exec };
|
||||
}
|
||||
|
||||
function makeService(redis: Redis) {
|
||||
const redisService = {
|
||||
getOrThrow: () => redis,
|
||||
} as unknown as RedisService;
|
||||
return new EmbeddingReindexProgressService(redisService);
|
||||
}
|
||||
|
||||
describe('get', () => {
|
||||
it('maps a valid hash to a ReindexProgress object', async () => {
|
||||
const { redis, hgetall } = makeRedis();
|
||||
hgetall.mockResolvedValue({ total: '478', done: '120', startedAt: '1000' });
|
||||
const service = makeService(redis);
|
||||
|
||||
await expect(service.get(WORKSPACE_ID)).resolves.toEqual({
|
||||
total: 478,
|
||||
done: 120,
|
||||
startedAt: 1000,
|
||||
});
|
||||
expect(hgetall).toHaveBeenCalledWith(KEY);
|
||||
});
|
||||
|
||||
it('returns null for an empty hash (no record)', async () => {
|
||||
const { redis, hgetall } = makeRedis();
|
||||
hgetall.mockResolvedValue({});
|
||||
await expect(makeService(redis).get(WORKSPACE_ID)).resolves.toBeNull();
|
||||
});
|
||||
|
||||
it('returns null when `total` is missing (partial record)', async () => {
|
||||
const { redis, hgetall } = makeRedis();
|
||||
hgetall.mockResolvedValue({ done: '5' });
|
||||
await expect(makeService(redis).get(WORKSPACE_ID)).resolves.toBeNull();
|
||||
});
|
||||
|
||||
it('returns null for a non-numeric total', async () => {
|
||||
const { redis, hgetall } = makeRedis();
|
||||
hgetall.mockResolvedValue({ total: 'abc', done: '1', startedAt: '1' });
|
||||
await expect(makeService(redis).get(WORKSPACE_ID)).resolves.toBeNull();
|
||||
});
|
||||
|
||||
it('returns null for a non-numeric done', async () => {
|
||||
const { redis, hgetall } = makeRedis();
|
||||
hgetall.mockResolvedValue({ total: '10', done: 'xyz', startedAt: '1' });
|
||||
await expect(makeService(redis).get(WORKSPACE_ID)).resolves.toBeNull();
|
||||
});
|
||||
|
||||
it('coerces a non-finite startedAt to 0', async () => {
|
||||
const { redis, hgetall } = makeRedis();
|
||||
hgetall.mockResolvedValue({ total: '10', done: '2', startedAt: 'nope' });
|
||||
await expect(makeService(redis).get(WORKSPACE_ID)).resolves.toEqual({
|
||||
total: 10,
|
||||
done: 2,
|
||||
startedAt: 0,
|
||||
});
|
||||
});
|
||||
|
||||
it('degrades to null when hgetall throws (degradation contract)', async () => {
|
||||
const { redis, hgetall } = makeRedis();
|
||||
hgetall.mockRejectedValue(new Error('redis down'));
|
||||
await expect(makeService(redis).get(WORKSPACE_ID)).resolves.toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe('start', () => {
|
||||
it('issues hset + expire on the workspace key', async () => {
|
||||
const { redis, multiObj } = makeRedis();
|
||||
await makeService(redis).start(WORKSPACE_ID, 478);
|
||||
|
||||
expect(multiObj.hset).toHaveBeenCalledWith(
|
||||
KEY,
|
||||
expect.objectContaining({ total: '478', done: '0' }),
|
||||
);
|
||||
expect(multiObj.expire).toHaveBeenCalledWith(KEY, expect.any(Number));
|
||||
expect(multiObj.exec).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('defaults the expire TTL to the full 1h record TTL', async () => {
|
||||
const { redis, multiObj } = makeRedis();
|
||||
await makeService(redis).start(WORKSPACE_ID, 478);
|
||||
// Default ttl = full record TTL (60 * 60) so a real run never expires
|
||||
// mid-flight before the worker refreshes it on each increment.
|
||||
expect(multiObj.expire).toHaveBeenCalledWith(KEY, 60 * 60);
|
||||
});
|
||||
|
||||
it('honours an explicit short ttlSeconds for the enqueue-time pre-seed (F10)', async () => {
|
||||
const { redis, multiObj } = makeRedis();
|
||||
// The reindex() pre-seed passes a short ttl so a phantom record left by a
|
||||
// de-duplicated enqueue expires in seconds, not after the full 1h TTL.
|
||||
await makeService(redis).start(WORKSPACE_ID, 478, 45);
|
||||
expect(multiObj.expire).toHaveBeenCalledWith(KEY, 45);
|
||||
});
|
||||
|
||||
it('swallows a thrown Redis error (best-effort)', async () => {
|
||||
const { redis } = makeRedis({
|
||||
execImpl: () => Promise.reject(new Error('redis down')),
|
||||
});
|
||||
await expect(
|
||||
makeService(redis).start(WORKSPACE_ID, 1),
|
||||
).resolves.toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('increment', () => {
|
||||
it('issues hincrby + expire on the workspace key', async () => {
|
||||
const { redis, multiObj } = makeRedis();
|
||||
await makeService(redis).increment(WORKSPACE_ID);
|
||||
|
||||
expect(multiObj.hincrby).toHaveBeenCalledWith(KEY, 'done', 1);
|
||||
expect(multiObj.expire).toHaveBeenCalledWith(KEY, expect.any(Number));
|
||||
expect(multiObj.exec).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('swallows a thrown Redis error (best-effort)', async () => {
|
||||
const { redis } = makeRedis({
|
||||
execImpl: () => Promise.reject(new Error('redis down')),
|
||||
});
|
||||
await expect(
|
||||
makeService(redis).increment(WORKSPACE_ID),
|
||||
).resolves.toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('clear', () => {
|
||||
it('deletes the workspace key', async () => {
|
||||
const { redis, del } = makeRedis();
|
||||
await makeService(redis).clear(WORKSPACE_ID);
|
||||
expect(del).toHaveBeenCalledWith(KEY);
|
||||
});
|
||||
|
||||
it('swallows a thrown Redis error (best-effort)', async () => {
|
||||
const { redis, del } = makeRedis();
|
||||
del.mockRejectedValue(new Error('redis down'));
|
||||
await expect(
|
||||
makeService(redis).clear(WORKSPACE_ID),
|
||||
).resolves.toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,162 @@
|
||||
import { Injectable, Logger } from '@nestjs/common';
|
||||
import { RedisService } from '@nestjs-labs/nestjs-ioredis';
|
||||
import type { Redis } from 'ioredis';
|
||||
|
||||
/**
|
||||
* Live progress of an in-flight workspace embeddings reindex run.
|
||||
* `total` is the number of pages the run will process, `done` how many it has
|
||||
* already processed (success OR handled failure), `startedAt` the epoch-ms the
|
||||
* record was created.
|
||||
*/
|
||||
export interface ReindexProgress {
|
||||
total: number;
|
||||
done: number;
|
||||
startedAt: number;
|
||||
}
|
||||
|
||||
/** Redis key namespace for the per-workspace reindex-progress record. */
|
||||
const KEY_PREFIX = 'ai:reindex:progress:';
|
||||
|
||||
/**
|
||||
* TTL (seconds) on the progress record so a crashed/aborted worker that never
|
||||
* reaches its `clear()` finally can still self-clean instead of leaving a stuck
|
||||
* "reindexing" state. Refreshed on every increment so a long run never expires
|
||||
* mid-flight; on a crash it disappears within TTL of the last processed page.
|
||||
*
|
||||
* INTENTIONALLY tied to WRITE progress (start/increment) only — never refreshed
|
||||
* on get(). Refreshing on read would keep a dead worker's record alive forever
|
||||
* as long as a client keeps polling (a permanently stuck reindexing:true). The
|
||||
* clear() in the worker's finally handles normal completion; a dead worker's
|
||||
* record expires after TTL, and the client's own poll cap stops polling anyway.
|
||||
*/
|
||||
const TTL_SECONDS = 60 * 60; // 1h
|
||||
|
||||
/**
|
||||
* Cluster-wide store for the live progress of a workspace embeddings reindex.
|
||||
*
|
||||
* The reindex runs in a BullMQ worker (AI_QUEUE) that may be a DIFFERENT process
|
||||
* than the API handling the settings-status GET, so the progress must live in
|
||||
* the shared Redis — we reuse the same global ioredis client (RedisService from
|
||||
* @nestjs-labs/nestjs-ioredis) that backs BullMQ and the other anti-abuse
|
||||
* limiters, adding NO new Redis config.
|
||||
*
|
||||
* Everything here is best-effort and COSMETIC: progress only drives the "Indexed
|
||||
* X of Y" counter while a reindex is running. Any Redis failure degrades to the
|
||||
* existing steady-state behaviour (the status falls back to the DB coverage
|
||||
* count), so reads fail to `null` and writes are swallowed — a reindex must
|
||||
* never break because progress reporting did.
|
||||
*
|
||||
* Stored as a Redis HASH so `done` can be bumped with an atomic HINCRBY (the
|
||||
* worker is the only writer of `done`, but HINCRBY also keeps us off a
|
||||
* read-modify-write race and preserves the other fields).
|
||||
*/
|
||||
@Injectable()
|
||||
export class EmbeddingReindexProgressService {
|
||||
private readonly logger = new Logger(EmbeddingReindexProgressService.name);
|
||||
private readonly redis: Redis;
|
||||
|
||||
constructor(redisService: RedisService) {
|
||||
this.redis = redisService.getOrThrow();
|
||||
}
|
||||
|
||||
private key(workspaceId: string): string {
|
||||
return KEY_PREFIX + workspaceId;
|
||||
}
|
||||
|
||||
/**
|
||||
* Begin (or reset) the progress record for a workspace: `total` pages, `done`
|
||||
* back to 0, `startedAt` now. Called twice for a run, BOTH with the real page
|
||||
* count (countEmbeddablePages) so the two totals coincide: once at reindex
|
||||
* enqueue time (so the very first status poll already reports done=0) and again
|
||||
* at the worker start (which re-asserts the same total and resets `done`).
|
||||
* Resets `done` to 0 so a re-trigger never inherits a stale count.
|
||||
*
|
||||
* `ttlSeconds` lets the caller pick the record's lifetime. The enqueue-time
|
||||
* pre-seed passes a SHORT ttl: if `aiQueue.add()` de-duplicates against a job
|
||||
* that is just finishing (its worker hasn't yet removed the job but already
|
||||
* ran its `clear()`), no new worker starts to clear this phantom seed, so a
|
||||
* short ttl lets it expire in seconds instead of sticking for the full TTL.
|
||||
* The worker's own `start()` at the begin of a real run overwrites this entry
|
||||
* and raises the ttl back to the default full TTL.
|
||||
*/
|
||||
async start(
|
||||
workspaceId: string,
|
||||
total: number,
|
||||
ttlSeconds: number = TTL_SECONDS,
|
||||
): Promise<void> {
|
||||
const key = this.key(workspaceId);
|
||||
try {
|
||||
await this.redis
|
||||
.multi()
|
||||
.hset(key, {
|
||||
total: String(total),
|
||||
done: '0',
|
||||
startedAt: String(Date.now()),
|
||||
})
|
||||
.expire(key, ttlSeconds)
|
||||
.exec();
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`reindex-progress start failed for workspace ${workspaceId}; ` +
|
||||
`progress reporting disabled for this run: ${(err as Error).message}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Bump the processed-page counter by one and refresh the TTL. Atomic and
|
||||
* best-effort: a missing key (cleared/expired) would be recreated with only
|
||||
* `done`, but `get()` treats a record without a numeric `total` as inactive,
|
||||
* so that partial state safely reads as "no active reindex".
|
||||
*/
|
||||
async increment(workspaceId: string): Promise<void> {
|
||||
const key = this.key(workspaceId);
|
||||
try {
|
||||
await this.redis.multi().hincrby(key, 'done', 1).expire(key, TTL_SECONDS).exec();
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`reindex-progress increment failed for workspace ${workspaceId}: ` +
|
||||
`${(err as Error).message}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Remove the progress record. Called in the worker's `finally` so a completed,
|
||||
* aborted, or unconfigured-early-return run never leaves a stuck record; the
|
||||
* status then falls back to the DB coverage count.
|
||||
*/
|
||||
async clear(workspaceId: string): Promise<void> {
|
||||
try {
|
||||
await this.redis.del(this.key(workspaceId));
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`reindex-progress clear failed for workspace ${workspaceId} ` +
|
||||
`(self-cleans via TTL): ${(err as Error).message}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Read the live progress, or `null` when no reindex is active (no record, an
|
||||
* expired record, or a partial record without a numeric `total`). On a Redis
|
||||
* error returns `null` so the status endpoint degrades to its DB count.
|
||||
*/
|
||||
async get(workspaceId: string): Promise<ReindexProgress | null> {
|
||||
try {
|
||||
const data = await this.redis.hgetall(this.key(workspaceId));
|
||||
if (!data || data.total === undefined) return null;
|
||||
const total = Number(data.total);
|
||||
const done = Number(data.done);
|
||||
const startedAt = Number(data.startedAt);
|
||||
if (!Number.isFinite(total) || !Number.isFinite(done)) return null;
|
||||
return { total, done, startedAt: Number.isFinite(startedAt) ? startedAt : 0 };
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`reindex-progress read failed for workspace ${workspaceId}; ` +
|
||||
`falling back to DB count: ${(err as Error).message}`,
|
||||
);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,18 +1,110 @@
|
||||
import { Readable } from 'stream';
|
||||
import { StorageService } from './storage.service';
|
||||
import type { StorageDriver } from './interfaces';
|
||||
|
||||
// Direct instantiation with a stub driver. The Test.createTestingModule form
|
||||
// failed to resolve the STORAGE_DRIVER_TOKEN at compile(); this smoke test only
|
||||
// needs the service to construct.
|
||||
describe('StorageService', () => {
|
||||
/**
|
||||
* StorageService is a thin facade over the injected StorageDriver: each public
|
||||
* method must forward to the driver with the SAME arguments and return/await the
|
||||
* driver's result unchanged (the read paths return it; the write paths await it).
|
||||
* A mock driver lets us assert that delegation exactly, with no real S3/disk IO.
|
||||
*/
|
||||
describe('StorageService delegation', () => {
|
||||
// Every driver method is a jest mock so we can assert call args + return passing.
|
||||
function buildDriver(): jest.Mocked<StorageDriver> {
|
||||
return {
|
||||
upload: jest.fn().mockResolvedValue(undefined),
|
||||
uploadStream: jest.fn().mockResolvedValue(undefined),
|
||||
copy: jest.fn().mockResolvedValue(undefined),
|
||||
read: jest.fn(),
|
||||
readStream: jest.fn(),
|
||||
readRangeStream: jest.fn(),
|
||||
exists: jest.fn(),
|
||||
getUrl: jest.fn(),
|
||||
getSignedUrl: jest.fn(),
|
||||
delete: jest.fn().mockResolvedValue(undefined),
|
||||
getDriver: jest.fn(),
|
||||
getDriverName: jest.fn(),
|
||||
getConfig: jest.fn(),
|
||||
} as unknown as jest.Mocked<StorageDriver>;
|
||||
}
|
||||
|
||||
let driver: jest.Mocked<StorageDriver>;
|
||||
let service: StorageService;
|
||||
|
||||
beforeEach(() => {
|
||||
service = new StorageService(
|
||||
{} as any, // storageDriver
|
||||
);
|
||||
driver = buildDriver();
|
||||
service = new StorageService(driver as unknown as StorageDriver);
|
||||
});
|
||||
|
||||
it('should be defined', () => {
|
||||
expect(service).toBeDefined();
|
||||
it('upload forwards path + content to the driver', async () => {
|
||||
const buf = Buffer.from('data');
|
||||
await service.upload('a/b.png', buf);
|
||||
expect(driver.upload).toHaveBeenCalledWith('a/b.png', buf);
|
||||
});
|
||||
|
||||
it('uploadStream forwards path, stream and options', async () => {
|
||||
const stream = Readable.from(['x']);
|
||||
await service.uploadStream('a/b.bin', stream, { recreateClient: true });
|
||||
expect(driver.uploadStream).toHaveBeenCalledWith('a/b.bin', stream, {
|
||||
recreateClient: true,
|
||||
});
|
||||
});
|
||||
|
||||
it('copy forwards both paths', async () => {
|
||||
await service.copy('from.txt', 'to.txt');
|
||||
expect(driver.copy).toHaveBeenCalledWith('from.txt', 'to.txt');
|
||||
});
|
||||
|
||||
it('read returns the driver buffer unchanged', async () => {
|
||||
const buf = Buffer.from('content');
|
||||
driver.read.mockResolvedValue(buf);
|
||||
await expect(service.read('f.txt')).resolves.toBe(buf);
|
||||
expect(driver.read).toHaveBeenCalledWith('f.txt');
|
||||
});
|
||||
|
||||
it('readStream returns the driver stream unchanged', async () => {
|
||||
const stream = Readable.from(['y']);
|
||||
driver.readStream.mockResolvedValue(stream);
|
||||
await expect(service.readStream('f.bin')).resolves.toBe(stream);
|
||||
expect(driver.readStream).toHaveBeenCalledWith('f.bin');
|
||||
});
|
||||
|
||||
it('readRangeStream forwards the range object and returns the stream', async () => {
|
||||
const stream = Readable.from(['z']);
|
||||
driver.readRangeStream.mockResolvedValue(stream);
|
||||
const range = { start: 0, end: 99 };
|
||||
await expect(service.readRangeStream('f.bin', range)).resolves.toBe(stream);
|
||||
expect(driver.readRangeStream).toHaveBeenCalledWith('f.bin', range);
|
||||
});
|
||||
|
||||
it('exists returns the driver boolean', async () => {
|
||||
driver.exists.mockResolvedValue(false);
|
||||
await expect(service.exists('missing')).resolves.toBe(false);
|
||||
expect(driver.exists).toHaveBeenCalledWith('missing');
|
||||
});
|
||||
|
||||
it('getSignedUrl forwards path + expiry and returns the signed url', async () => {
|
||||
driver.getSignedUrl.mockResolvedValue('https://signed/url');
|
||||
await expect(service.getSignedUrl('f.png', 600)).resolves.toBe(
|
||||
'https://signed/url',
|
||||
);
|
||||
expect(driver.getSignedUrl).toHaveBeenCalledWith('f.png', 600);
|
||||
});
|
||||
|
||||
it('getUrl returns the driver url synchronously', () => {
|
||||
driver.getUrl.mockReturnValue('https://cdn/f.png');
|
||||
expect(service.getUrl('f.png')).toBe('https://cdn/f.png');
|
||||
expect(driver.getUrl).toHaveBeenCalledWith('f.png');
|
||||
});
|
||||
|
||||
it('delete forwards the path', async () => {
|
||||
await service.delete('old.txt');
|
||||
expect(driver.delete).toHaveBeenCalledWith('old.txt');
|
||||
});
|
||||
|
||||
it('getDriverName returns the driver name', () => {
|
||||
driver.getDriverName.mockReturnValue('s3');
|
||||
expect(service.getDriverName()).toBe('s3');
|
||||
expect(driver.getDriverName).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -10,7 +10,6 @@ import {
|
||||
PAGE_TEMPLATE_THROTTLER,
|
||||
PUBLIC_SHARE_AI_THROTTLER,
|
||||
} from './throttler-names';
|
||||
import Redis from 'ioredis';
|
||||
|
||||
@Module({
|
||||
imports: [
|
||||
@@ -32,16 +31,18 @@ import Redis from 'ioredis';
|
||||
{ name: PUBLIC_SHARE_AI_THROTTLER, ttl: 60_000, limit: 5 },
|
||||
],
|
||||
errorMessage: 'Too many requests',
|
||||
storage: new ThrottlerStorageRedisService(
|
||||
new Redis({
|
||||
host: redisConfig.host,
|
||||
port: redisConfig.port,
|
||||
password: redisConfig.password,
|
||||
db: redisConfig.db,
|
||||
family: redisConfig.family,
|
||||
keyPrefix: 'throttle:',
|
||||
}),
|
||||
),
|
||||
// Pass ioredis options (not a pre-built Redis instance) so
|
||||
// ThrottlerStorageRedisService owns the connection and disconnects it
|
||||
// in its onModuleDestroy. Passing an instance leaves disconnectRequired
|
||||
// false, so the socket would leak on shutdown (e2e jest never exits).
|
||||
storage: new ThrottlerStorageRedisService({
|
||||
host: redisConfig.host,
|
||||
port: redisConfig.port,
|
||||
password: redisConfig.password,
|
||||
db: redisConfig.db,
|
||||
family: redisConfig.family,
|
||||
keyPrefix: 'throttle:',
|
||||
}),
|
||||
};
|
||||
},
|
||||
inject: [EnvironmentService],
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import { Logger } from '@nestjs/common';
|
||||
import { IoAdapter } from '@nestjs/platform-socket.io';
|
||||
import { ServerOptions } from 'socket.io';
|
||||
import { createAdapter } from '@socket.io/redis-adapter';
|
||||
@@ -9,8 +10,11 @@ import {
|
||||
} from '../../common/helpers';
|
||||
|
||||
export class WsRedisIoAdapter extends IoAdapter {
|
||||
private readonly logger = new Logger(WsRedisIoAdapter.name);
|
||||
private adapterConstructor: ReturnType<typeof createAdapter>;
|
||||
private redisConfig: RedisConfig;
|
||||
private pubClient: Redis;
|
||||
private subClient: Redis;
|
||||
|
||||
async connectToRedis(): Promise<void> {
|
||||
this.redisConfig = parseRedisUrl(process.env.REDIS_URL);
|
||||
@@ -23,8 +27,13 @@ export class WsRedisIoAdapter extends IoAdapter {
|
||||
const pubClient = new Redis(process.env.REDIS_URL, options);
|
||||
const subClient = new Redis(process.env.REDIS_URL, options);
|
||||
|
||||
pubClient.on('error', (err) => () => {});
|
||||
subClient.on('error', (err) => () => {});
|
||||
pubClient.on('error', (err) => this.logger.error('socket.io redis pub client error', err));
|
||||
subClient.on('error', (err) => this.logger.error('socket.io redis sub client error', err));
|
||||
|
||||
// Hold references so the pub/sub connections can be torn down on shutdown
|
||||
// (see dispose()); otherwise these ioredis sockets leak as active handles.
|
||||
this.pubClient = pubClient;
|
||||
this.subClient = subClient;
|
||||
|
||||
this.adapterConstructor = createAdapter(pubClient, subClient);
|
||||
}
|
||||
@@ -34,4 +43,26 @@ export class WsRedisIoAdapter extends IoAdapter {
|
||||
server.adapter(this.adapterConstructor);
|
||||
return server;
|
||||
}
|
||||
|
||||
/**
|
||||
* Called once by Nest's SocketModule during application shutdown, after every
|
||||
* socket.io server has been closed. The @socket.io/redis-adapter never owns
|
||||
* the lifecycle of the ioredis pub/sub clients it is handed, so we close them
|
||||
* here to avoid leaking their TCP handles on shutdown (see issue #255).
|
||||
*
|
||||
* Uses disconnect(false) to mirror the sibling pub/sub pair in
|
||||
* collaboration/extensions/redis-sync (redis-sync.extension.ts onDestroy):
|
||||
* an immediate close with no graceful QUIT round-trip and no auto-reconnect,
|
||||
* which is what we want for idle adapter clients during teardown.
|
||||
*/
|
||||
async dispose(): Promise<void> {
|
||||
await super.dispose();
|
||||
|
||||
// dispose() is invoked once per shutdown; null the refs so a second call
|
||||
// (or any post-shutdown path) cannot act on already-closed clients.
|
||||
this.pubClient?.disconnect(false);
|
||||
this.subClient?.disconnect(false);
|
||||
this.pubClient = undefined;
|
||||
this.subClient = undefined;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user