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
@@ -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],
|
||||
|
||||
Reference in New Issue
Block a user