diff --git a/CHANGELOG.md b/CHANGELOG.md index 47ce54f5..3d8958d4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -364,6 +364,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 before retry") so the agent re-reads instead of blind-retrying, and a still-connecting session is no longer picked as an idle eviction victim by a parallel acquire. (#494) +- **A long AI chat no longer bricks on the model's context window, and each turn + stops re-persisting the whole tool-output history.** Tool outputs are now + stored ONCE, in `metadata.parts`; the `tool_calls` trace keeps only per-step + outcome flags (a v2 trace shape), ending the O(N²) write amplification that + re-wrote every prior output on every step (measured on a live Postgres via the + `pg_current_wal_lsn()` delta: the trace column shrank ~3200×, the full + assistant row ~51%). The persisted record is unchanged in content — the full + history still lives in `metadata.parts`. At REPLAY time only, the history sent + to the provider is now bounded by a deterministic, prompt-cache-friendly token + budget: `floor(0.7 × chatContextWindow)` when a window is configured (no cap — + anti-brick protection, not a cost limiter), a flat 100k fallback for installs + with no window set (exactly the ones that hit terminal overflow), or off when + the window is explicitly `0`. Trimming truncates old tool outputs first, then + mechanically collapses the oldest turns, always keeping the recent turns full + and the tool-call/result pairing balanced. A provider context-overflow 400 is + now classified and used as a reactive signal: the row is stamped so the NEXT + turn re-trims aggressively (0.5×), which un-bricks a chat that just 400'd. The + client token badge and the server budgeter now share one estimator (new + `@docmost/token-estimate` package) so they can never diverge. Deferred-tool + activation is also cached in the chat metadata to avoid re-resolving it each + turn. (#490) - **A chat with one malformed message part no longer 500s on every turn, and a failed send no longer duplicates the user's message.** Incoming client parts are now whitelisted to `text` (a forged tool-result part can no longer reach diff --git a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts index 7c0e4273..5d4ff1b6 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts @@ -36,6 +36,7 @@ import { import type { AiChatMessage, Workspace } from '@docmost/db/types/entity.types'; import { buildSystemPrompt } from './ai-chat.prompt'; import type { McpClientsService } from './external-mcp/mcp-clients.service'; +import { resolveEffectiveReplayThreshold } from './history-budget'; /** * Unit tests for compactToolOutput: the pure helper that shrinks tool outputs @@ -576,6 +577,27 @@ describe('lastAssistantReplayOverflow', () => { ).toBe(false); expect(lastAssistantReplayOverflow([])).toBe(false); }); + + // #490 reactive recovery: a prior turn stamped `replayOverflow` must make the + // NEXT turn's effective budget the AGGRESSIVE 0.5x cut — that harder trim is + // what un-bricks a chat that just 400'd on the context window. This exercises + // the exact wiring the service uses: read the stamp, then scale the threshold. + it('#490: a prior replayOverflow drives the next turn to the 0.5x aggressive budget', () => { + const history = [ + row('assistant', { replayOverflow: true }), + row('user', null), + ]; + const priorOverflowed = lastAssistantReplayOverflow(history); + expect(priorOverflowed).toBe(true); + // Base budget 100k -> aggressive recovery halves it to 50k this turn. + expect(resolveEffectiveReplayThreshold(100_000, priorOverflowed)).toBe(50_000); + // Odd base floors, not rounds. + expect(resolveEffectiveReplayThreshold(99_999, true)).toBe(49_999); + // No prior overflow -> the base budget is used verbatim (no aggressive cut). + expect(resolveEffectiveReplayThreshold(100_000, false)).toBe(100_000); + // An explicit off-switch (null) is never overridden, even on recovery. + expect(resolveEffectiveReplayThreshold(null, true)).toBeNull(); + }); }); describe('rowToUiMessage', () => { diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index e1626d81..f1cda9b4 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -57,9 +57,9 @@ import { import { roleModelOverride } from './roles/role-model-config'; import { resolveReplayBudget, + resolveEffectiveReplayThreshold, isContextOverflowError, trimHistoryForReplay, - REPLAY_AGGRESSIVE_FRACTION, } from './history-budget'; import { startSseHeartbeat, @@ -1197,12 +1197,10 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // overflowing turn produced no usage signal, so a normal-threshold trim may // not shrink enough to fit. This is what un-bricks a chat that just 400'd. const priorOverflowed = lastAssistantReplayOverflow(oldHistory); - const effectiveThreshold = - priorOverflowed && replayBudget.thresholdTokens != null - ? Math.floor( - replayBudget.thresholdTokens * REPLAY_AGGRESSIVE_FRACTION, - ) - : replayBudget.thresholdTokens; + const effectiveThreshold = resolveEffectiveReplayThreshold( + replayBudget.thresholdTokens, + priorOverflowed, + ); if (priorOverflowed) { this.logger.warn( `AI chat (chat ${chatId}): previous turn hit context overflow; ` + diff --git a/apps/server/src/core/ai-chat/ai-chat.write-volume.spec.ts b/apps/server/src/core/ai-chat/ai-chat.write-volume.int-spec.ts similarity index 100% rename from apps/server/src/core/ai-chat/ai-chat.write-volume.spec.ts rename to apps/server/src/core/ai-chat/ai-chat.write-volume.int-spec.ts diff --git a/apps/server/src/core/ai-chat/history-budget.spec.ts b/apps/server/src/core/ai-chat/history-budget.spec.ts index 16f6fd76..5f1fb72b 100644 --- a/apps/server/src/core/ai-chat/history-budget.spec.ts +++ b/apps/server/src/core/ai-chat/history-budget.spec.ts @@ -10,15 +10,15 @@ import { } from './history-budget'; describe('resolveReplayBudget', () => { - it('uses min(default, 0.7 x window) for a configured window', () => { - // 0.7 x 60k = 42k < 100k + it('uses floor(0.7 x window) for a configured window (no cap)', () => { + // 0.7 x 60k = 42k expect(resolveReplayBudget(60_000)).toEqual({ thresholdTokens: 42_000, usedDefault: false, }); - // 0.7 x 1M = 700k, capped to the 100k default + // 0.7 x 1M = 700k — NOT capped (anti-brick vs the window, not a cost limiter). expect(resolveReplayBudget(1_000_000)).toEqual({ - thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS, + thresholdTokens: 700_000, usedDefault: false, }); }); diff --git a/apps/server/src/core/ai-chat/history-budget.ts b/apps/server/src/core/ai-chat/history-budget.ts index fb45702f..38cb6465 100644 --- a/apps/server/src/core/ai-chat/history-budget.ts +++ b/apps/server/src/core/ai-chat/history-budget.ts @@ -51,7 +51,10 @@ export interface ReplayBudget { /** * Resolve the replay budget from the RAW stored `chatContextWindow` (text/number). - * - a positive value -> `min(default, floor(fraction × window))` + * - a positive value -> `floor(fraction × window)` (NO cap — the budgeter is + * anti-brick protection against the window itself, not a cost/economy limiter, + * exactly as the codebase already treats maxOutputTokens; the reactive branch + * still guarantees anti-brick regardless of how high this budget is) * - explicit `0` -> OFF (admin opt-out; `null` threshold) * - unset/empty/invalid-> the flat default (still protects — the installations * that hit terminal overflow are exactly the ones that never set a window) @@ -76,14 +79,30 @@ export function resolveReplayBudget(rawContextWindow: unknown): ReplayBudget { return { thresholdTokens: null, usedDefault: false }; } return { - thresholdTokens: Math.min( - REPLAY_BUDGET_DEFAULT_TOKENS, - Math.floor(REPLAY_BUDGET_WINDOW_FRACTION * n), - ), + thresholdTokens: Math.floor(REPLAY_BUDGET_WINDOW_FRACTION * n), usedDefault: false, }; } +/** + * The effective replay threshold for THIS turn, given the base budget and whether + * the PREVIOUS turn hit a context-overflow 400 (the reactive-recovery signal, + * `metadata.replayOverflow`). On recovery the base budget is scaled down by + * {@link REPLAY_AGGRESSIVE_FRACTION}: the overflowing turn produced no usage + * signal, so the preventive estimate under-counted and a normal-threshold trim may + * not shrink enough to fit — this harder cut is what un-bricks the chat. + * + * A `null` base budget (trimming OFF) is passed through unchanged: an explicit + * off-switch is never overridden by the recovery path. + */ +export function resolveEffectiveReplayThreshold( + thresholdTokens: number | null, + priorOverflowed: boolean, +): number | null { + if (!priorOverflowed || thresholdTokens == null) return thresholdTokens; + return Math.floor(thresholdTokens * REPLAY_AGGRESSIVE_FRACTION); +} + /** * True when a provider error is a CONTEXT-OVERFLOW rejection (the prompt exceeds * the model's window). Providers surface this as an HTTP 400 with a recognizable diff --git a/apps/server/test/jest-integration.json b/apps/server/test/jest-integration.json index 30f62397..cd3633c5 100644 --- a/apps/server/test/jest-integration.json +++ b/apps/server/test/jest-integration.json @@ -22,6 +22,7 @@ "^@docmost/db/(.*)$": "/src/database/$1", "^@docmost/transactional/(.*)$": "/src/integrations/transactional/$1", "^@docmost/ee/(.*)$": "/src/ee/$1", + "^@docmost/token-estimate$": "/../../packages/token-estimate/src/index.ts", "^src/(.*)$": "/src/$1" } } diff --git a/packages/mcp/test/mock/pagination-cursor.test.mjs b/packages/mcp/test/mock/pagination-cursor.test.mjs index efc997cd..39ce55b8 100644 --- a/packages/mcp/test/mock/pagination-cursor.test.mjs +++ b/packages/mcp/test/mock/pagination-cursor.test.mjs @@ -510,3 +510,71 @@ test("checkNewComments fetches pages concurrently (bounded) and preserves order" "result order matches the enumeration order regardless of finish order", ); }); + +// ----------------------------------------------------------------------------- +// 7) checkNewComments partial failure (#490): the concurrent scan is resilient — +// if ONE page's /comments fetch rejects (deleted mid-scan, a transient 500), +// that page is skipped and the WHOLE scan still resolves with every other +// page's fresh comments. A single failing fetch must never reject the batch +// (Promise.all in mapWithConcurrency would otherwise abort all of it) nor +// corrupt the deterministic order of the pages that DID succeed. +// ----------------------------------------------------------------------------- +test("checkNewComments skips a page whose fetch fails and still reports the rest", async () => { + const NODES = [{ id: "parent", title: "Parent", parentPageId: null, hasChildren: true }]; + for (let i = 0; i < 5; i++) { + NODES.push({ id: `k${i}`, title: `Kid ${i}`, parentPageId: "parent", hasChildren: false }); + } + // The one page whose comment fetch blows up (500 -> listComments rejects). + const FAILING = "k2"; + + const { baseURL } = await spawn(async (req, res) => { + const raw = await readBody(req); + if (handleLogin(req, res)) return; + if (req.url === "/api/pages/tree") { + sendJson(res, 200, { success: true, data: { items: NODES } }); + return; + } + if (req.url === "/api/comments") { + const body = JSON.parse(raw || "{}"); + if (body.pageId === FAILING) { + // A transient server error on exactly one page's fetch. + sendJson(res, 500, { success: false, message: "boom" }); + return; + } + sendJson(res, 200, { + success: true, + data: { + items: [ + { id: `c-${body.pageId}`, createdAt: "2030-01-01T00:00:00.000Z", content: null }, + ], + meta: { nextCursor: null }, + }, + }); + return; + } + sendJson(res, 404, {}); + }); + + const client = new DocmostClient(baseURL, "user@example.com", "pw"); + // Must RESOLVE (not reject) despite one page's fetch failing. + const result = await client.checkNewComments( + "space-1", + "2020-01-01T00:00:00.000Z", + "parent", + ); + + // Every page in scope was still scanned (the failing one counts as checked). + assert.equal(result.checkedPages, 6, "all pages scanned incl. the failing one"); + // The failing page contributes nothing; the other 5 each report one comment. + assert.equal(result.pagesWithNewComments, 5, "the failing page is dropped"); + assert.equal(result.totalNewComments, 5, "only the succeeding pages' comments"); + const reportedIds = result.comments.map((r) => r.pageId); + assert.ok(!reportedIds.includes(FAILING), "the failing page is absent from results"); + // Order of the survivors is still the deterministic enumeration order (the hole + // left by the failing page is closed without reordering the rest). + assert.deepEqual( + reportedIds, + ["parent", "k0", "k1", "k3", "k4"], + "survivors keep enumeration order with the failing page removed", + ); +});