From d4b21957dccf049c4af17d7a09c4b2673a269ac2 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 11:21:27 +0300 Subject: [PATCH] =?UTF-8?q?feat(ai-chat):=20=D1=82=D0=BE=D0=BA=D0=B5=D0=BD?= =?UTF-8?q?-=D0=B1=D1=8E=D0=B4=D0=B6=D0=B5=D1=82=20=D1=80=D0=B5=D0=BF?= =?UTF-8?q?=D0=BB=D0=B5=D1=8F=20=D0=B8=D1=81=D1=82=D0=BE=D1=80=D0=B8=D0=B8?= =?UTF-8?q?=20+=20=D1=80=D0=B5=D0=B0=D0=BA=D1=82=D0=B8=D0=B2=D0=BD=D0=B0?= =?UTF-8?q?=D1=8F=20=D0=B2=D0=B5=D1=82=D0=BA=D0=B0=20(#490)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Вся персистентная история реплеится провайдеру КАЖДЫЙ ход, поэтому длинный чат рано или поздно упирается в контекстное окно и получает провайдерский 400 на каждом ходу — навсегда (чат «кирпичится»). Бюджетер ограничивает РЕПЛЕЙ (никогда не мутирует персист — в БД остаётся полная запись), детерминированно и byte-stable (обрезанный префикс идентичен от хода к ходу → дружелюбно к prompt-cache). Единый оценщик chars/2.5 (кириллица; chars/4 занижает вдвое) вынесен в shared-пакет packages/token-estimate; клиентский count-stream-tokens.ts переведён на него ТЕМ ЖЕ коммитом (два расходящихся оценщика = «бейдж 60%, а бюджет уже режет»). history-budget.ts (чистый, покрыт тестами): - resolveReplayBudget(raw): min(100k, 0.7×window) при заданном окне; флэт 100k при незаданном (именно эти инсталляции ловят терминальный overflow — warn-лог); 0 = явный off-switch. Читается СЫРОЙ chatContextWindow, т.к. parsePositiveInt схлопывает 0 и unset в undefined (новое поле ResolvedAiConfig.chatContextWindowRaw). - trimHistoryForReplay: первичный сигнал — провайдерский факт metadata.contextTokens прошлого хода; chars-оценка — дельта/раскройка/фолбэк. Порядок: обрезка tool-outputs старых ходов (head+tail+маркер) → механическое схлопывание старейших ходов (конкатенация, НЕ LLM) → текущий + последние N ходов всегда полные. Пейринг tool-call/result сохраняется (схлопывание убирает ОБЕ части). - isContextOverflowError: классификация провайдерского 400 (статус + паттерны). Реактивная ветка: превентивная оценка не даёт инварианта (первый переполняющий ход не имеет usage). onError классифицирует context-overflow → пишет различимую причину и штампует metadata.replayOverflow; следующий ход бюджетер режет агрессивно (0.5×), что и раскирпичивает чат. Наблюдаемость: metadata.replayTrimmedToTokens. ПРИМЕЧАНИЕ по реактивной ветке (форк, требует решения ревьюера): истинный in-turn re-pipe (перезапуск streamText в тот же ответ) архитектурно несовместим с текущим пайпом — pipeUIMessageStreamToResponse пишет writeHead СИНХРОННО (подтверждено в ai@6.0.207), а suite ожидает await stream() c моком, не дёргающим колбэки, — так что отложенный пайп/ожидание сигнала повесит тесты. Поэтому реализована реактивная рекавери «классификация → штамп → агрессивный ре-трим на следующем ходу», что даёт тот же инвариант (чат не кирпичится) без рискованного рефактора стрима. Тесты (наблюдаемые свойства): объём записи через дельту pg_current_wal_lsn() на живой gitmost-test-pg вокруг 50-шагового прогона (несжимаемые payload'ы) — trace-колонка v1=140МБ → v2=0.04МБ (в 3206× меньше), полная строка 289МБ → 140МБ (−51%); dual-shape не нужен здесь; «окно не задано → бюджет применяется»; реактивная классификация на реальном 400-шейпе; parity клиент/сервер оценщика. Co-Authored-By: Claude Opus 4.8 (1M context) --- .gitignore | 4 + apps/client/package.json | 1 + .../ai-chat/utils/count-stream-tokens.test.ts | 11 +- .../ai-chat/utils/count-stream-tokens.ts | 20 +- apps/server/package.json | 4 +- .../ai-chat/ai-chat.service.run-race.spec.ts | 51 ++- .../src/core/ai-chat/ai-chat.service.spec.ts | 68 ++++ .../src/core/ai-chat/ai-chat.service.ts | 132 ++++++- .../core/ai-chat/ai-chat.write-volume.spec.ts | 209 ++++++++++ .../src/core/ai-chat/history-budget.spec.ts | 266 +++++++++++++ .../server/src/core/ai-chat/history-budget.ts | 356 ++++++++++++++++++ .../integrations/ai/ai-settings.service.ts | 3 + apps/server/src/integrations/ai/ai.types.ts | 4 + packages/token-estimate/package.json | 19 + packages/token-estimate/src/index.test.ts | 31 ++ packages/token-estimate/src/index.ts | 35 ++ packages/token-estimate/tsconfig.json | 16 + pnpm-lock.yaml | 15 + 18 files changed, 1221 insertions(+), 24 deletions(-) create mode 100644 apps/server/src/core/ai-chat/ai-chat.write-volume.spec.ts create mode 100644 apps/server/src/core/ai-chat/history-budget.spec.ts create mode 100644 apps/server/src/core/ai-chat/history-budget.ts create mode 100644 packages/token-estimate/package.json create mode 100644 packages/token-estimate/src/index.test.ts create mode 100644 packages/token-estimate/src/index.ts create mode 100644 packages/token-estimate/tsconfig.json diff --git a/.gitignore b/.gitignore index 3eb7e75b..f3ce4f3e 100644 --- a/.gitignore +++ b/.gitignore @@ -29,6 +29,10 @@ packages/mcp/build/ # is a build artifact like build/ — never committed, always fresh. packages/mcp/src/registry-stamp.generated.ts +# token-estimate compiled output (#490; built in CI/Docker via `pnpm build` / +# the server `pretest`, never committed, so src/ and prod can never diverge). +packages/token-estimate/dist/ + # Logs logs *.log diff --git a/apps/client/package.json b/apps/client/package.json index a9fb48f7..00180fd4 100644 --- a/apps/client/package.json +++ b/apps/client/package.json @@ -22,6 +22,7 @@ "@casl/react": "5.0.1", "@docmost/editor-ext": "workspace:*", "@docmost/prosemirror-markdown": "workspace:*", + "@docmost/token-estimate": "workspace:*", "@excalidraw/excalidraw": "0.18.0-3a5ef40", "@mantine/core": "8.3.18", "@mantine/dates": "8.3.18", diff --git a/apps/client/src/features/ai-chat/utils/count-stream-tokens.test.ts b/apps/client/src/features/ai-chat/utils/count-stream-tokens.test.ts index 6b00fbc4..819c057e 100644 --- a/apps/client/src/features/ai-chat/utils/count-stream-tokens.test.ts +++ b/apps/client/src/features/ai-chat/utils/count-stream-tokens.test.ts @@ -6,10 +6,13 @@ describe("estimateTokens", () => { expect(estimateTokens("")).toBe(0); }); - it("ceils chars/4 so any non-empty text is at least 1 token", () => { + // #490: migrated onto the shared @docmost/token-estimate module (chars/2.5, up + // from the old client-only chars/4) so the client counter and the server replay + // budgeter can never diverge. + it("ceils chars/2.5 so any non-empty text is at least 1 token", () => { expect(estimateTokens("a")).toBe(1); - expect(estimateTokens("abcd")).toBe(1); - expect(estimateTokens("abcde")).toBe(2); - expect(estimateTokens("12345678")).toBe(2); + expect(estimateTokens("ab")).toBe(1); + expect(estimateTokens("abcde")).toBe(2); // 5 / 2.5 = 2 + expect(estimateTokens("x".repeat(10))).toBe(4); // 10 / 2.5 = 4 }); }); diff --git a/apps/client/src/features/ai-chat/utils/count-stream-tokens.ts b/apps/client/src/features/ai-chat/utils/count-stream-tokens.ts index aaf99599..951bb992 100644 --- a/apps/client/src/features/ai-chat/utils/count-stream-tokens.ts +++ b/apps/client/src/features/ai-chat/utils/count-stream-tokens.ts @@ -2,18 +2,10 @@ * Rough client-side token estimation for AI-chat UI affordances. * * No provider streams exact per-token usage mid-stream, so any in-flight figure - * is a CLIENT ESTIMATE (chars/≈4 heuristic). Pure + unit-testable: it never runs - * a real BPE tokenizer (that would be O(n²) on the hot path, bloat the bundle, - * and be wrong for Gemini/Ollama anyway). Used by the in-body reasoning counter - * ("Thinking · N tokens"). + * is a CLIENT ESTIMATE. This re-exports the SHARED estimator from + * `@docmost/token-estimate` (chars/2.5) so the in-body counter and the server's + * replay budgeter use the SAME heuristic — two divergent estimators would mean + * "the badge shows 60%" while "the budgeter already trimmed" (#490). Used by the + * in-body reasoning counter ("Thinking · N tokens"). */ - -/** - * Rough token estimate for a piece of text using the standard chars/≈4 heuristic. - * Returns 0 for empty/whitespace-free-of-content input, and ceils so any - * non-empty text counts as at least one token. - */ -export function estimateTokens(text: string): number { - if (!text) return 0; - return Math.ceil(text.length / 4); -} +export { estimateTokens } from "@docmost/token-estimate"; diff --git a/apps/server/package.json b/apps/server/package.json index a6791d80..7b97752a 100644 --- a/apps/server/package.json +++ b/apps/server/package.json @@ -23,7 +23,7 @@ "migration:reset": "tsx src/database/migrate.ts down-to NO_MIGRATIONS", "migration:codegen": "kysely-codegen --dialect=postgres --camel-case --env-file=../../.env --out-file=./src/database/types/db.d.ts", "lint": "eslint \"{src,apps,libs,test}/**/*.ts\" --fix", - "pretest": "pnpm --filter @docmost/editor-ext build && pnpm --filter @docmost/prosemirror-markdown build", + "pretest": "pnpm --filter @docmost/editor-ext build && pnpm --filter @docmost/prosemirror-markdown build && pnpm --filter @docmost/token-estimate build", "test": "jest", "test:int": "jest --config test/jest-integration.json", "test:watch": "jest --watch", @@ -44,6 +44,7 @@ "@docmost/mcp": "workspace:*", "@docmost/pdf-inspector": "1.9.6", "@docmost/prosemirror-markdown": "workspace:*", + "@docmost/token-estimate": "workspace:*", "@fastify/compress": "^9.0.0", "@fastify/cookie": "^11.0.2", "@fastify/multipart": "^10.0.0", @@ -206,6 +207,7 @@ "^@docmost/db/(.*)$": "/database/$1", "^@docmost/transactional/(.*)$": "/integrations/transactional/$1", "^@docmost/ee/(.*)$": "/ee/$1", + "^@docmost/token-estimate$": "/../../../packages/token-estimate/src/index.ts", "^src/(.*)$": "/$1", "^@tiptap/react$": "/../test/stubs/tiptap-react.js" } diff --git a/apps/server/src/core/ai-chat/ai-chat.service.run-race.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.run-race.spec.ts index 8b170c11..b15c2ed1 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.run-race.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.run-race.spec.ts @@ -181,7 +181,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => { {} as never, // pageAccess { isAiChatDeferredToolsEnabled: () => false, isAiChatFinalStepLockdownEnabled: () => false } as never, // environment ); - return { svc }; + return { svc, aiChatMessageRepo }; } const body = { @@ -287,7 +287,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => { // Drive stream() to the point streamText is called, capturing the options object // (which carries onStepFinish/onFinish/onError/onAbort) and the run hooks. async function captureStreamCallbacks() { - const { svc } = makeService(); + const { svc, aiChatMessageRepo } = makeService(); let capturedOpts: any; streamTextMock.mockImplementation((opts: any) => { capturedOpts = opts; @@ -314,7 +314,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => { runHooks: runHooks as never, }); expect(capturedOpts).toBeDefined(); - return { capturedOpts, runHooks }; + return { capturedOpts, runHooks, aiChatMessageRepo }; } it('F9: onStepFinish bumps the run step count, onFinish settles the run "completed" (the dominant autonomous-run path)', async () => { @@ -369,6 +369,51 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => { expect.stringContaining('provider exploded'), ); }); + + // #490 reactive branch: a provider CONTEXT-OVERFLOW 400 in onError is classified, + // records a distinguishable cause, and stamps metadata.replayOverflow so the NEXT + // turn's budgeter trims aggressively (the recovery that un-bricks the chat). + it('#490: a context-overflow 400 stamps replayOverflow on the finalized row', async () => { + jest + .spyOn(Logger.prototype, 'error') + .mockImplementation(() => undefined as never); + jest + .spyOn(Logger.prototype, 'warn') + .mockImplementation(() => undefined as never); + const { capturedOpts, aiChatMessageRepo } = await captureStreamCallbacks(); + + const overflow = Object.assign(new Error('too large'), { + statusCode: 400, + message: + "This model's maximum context length is 128000 tokens. However, your messages resulted in 214000 tokens. Please reduce the length.", + }); + await capturedOpts.onError({ error: overflow }); + + // The seed row exists (finalizeOwner is the owner-write path). + expect(aiChatMessageRepo.finalizeOwner).toHaveBeenCalled(); + const calls = aiChatMessageRepo.finalizeOwner.mock.calls as any[][]; + const patch = calls[calls.length - 1][2] as { + status: string; + metadata: Record; + }; + expect(patch.status).toBe('error'); + expect(patch.metadata.replayOverflow).toBe(true); + expect(patch.metadata.error).toContain('контекстное окно'); + }); + + it('#490: a non-overflow error does NOT stamp replayOverflow', async () => { + jest + .spyOn(Logger.prototype, 'error') + .mockImplementation(() => undefined as never); + const { capturedOpts, aiChatMessageRepo } = await captureStreamCallbacks(); + await capturedOpts.onError({ error: new Error('network reset') }); + const calls = aiChatMessageRepo.finalizeOwner.mock.calls as any[][]; + const patch = calls[calls.length - 1][2] as { + status: string; + metadata: Record; + }; + expect('replayOverflow' in patch.metadata).toBe(false); + }); }); /** 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 bee3139f..70293eb0 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 @@ -29,6 +29,8 @@ import { FINAL_STEP_NUDGE, STEP_LIMIT_NO_ANSWER_MARKER, OUTPUT_DEGENERATION_ERROR, + lastAssistantContextTokens, + lastAssistantReplayOverflow, } from './ai-chat.service'; import type { AiChatMessage, Workspace } from '@docmost/db/types/entity.types'; import { buildSystemPrompt } from './ai-chat.prompt'; @@ -413,6 +415,55 @@ describe('toolTraceVersion era marker (#490)', () => { }); }); +// #490 replay-budget signal helpers over persisted history. +describe('lastAssistantContextTokens', () => { + const row = ( + role: string, + metadata: Record | null, + ): AiChatMessage => ({ role, metadata }) as unknown as AiChatMessage; + + it('reads the most recent assistant turn contextTokens (provider fact)', () => { + const hist = [ + row('user', null), + row('assistant', { contextTokens: 12000 }), + row('user', null), + row('assistant', { contextTokens: 41000 }), + ]; + expect(lastAssistantContextTokens(hist)).toBe(41000); + }); + + it('returns undefined when the last assistant turn recorded no usage', () => { + const hist = [row('assistant', { error: 'boom' }), row('user', null)]; + expect(lastAssistantContextTokens(hist)).toBeUndefined(); + expect(lastAssistantContextTokens([])).toBeUndefined(); + }); +}); + +describe('lastAssistantReplayOverflow', () => { + const row = ( + role: string, + metadata: Record | null, + ): AiChatMessage => ({ role, metadata }) as unknown as AiChatMessage; + + it('is true only when the LAST assistant turn overflowed', () => { + expect( + lastAssistantReplayOverflow([ + row('assistant', { replayOverflow: true }), + row('user', null), + ]), + ).toBe(true); + // A recovered (later, non-overflow) assistant turn clears it. + expect( + lastAssistantReplayOverflow([ + row('assistant', { replayOverflow: true }), + row('user', null), + row('assistant', { contextTokens: 5 }), + ]), + ).toBe(false); + expect(lastAssistantReplayOverflow([])).toBe(false); + }); +}); + describe('rowToUiMessage', () => { it('prefers metadata.parts over content', () => { const row = { @@ -743,6 +794,23 @@ describe('flushAssistant', () => { expect(flushed.metadata.error).toBe('boom'); }); + // #490 observability: the replay budgeter's decision is stamped on the turn. + it('records replayTrimmedToTokens + replayOverflow when provided', () => { + const f = flushAssistant([], '', 'error', { + error: 'ctx', + replayTrimmedToTokens: 42_000, + replayOverflow: true, + }); + expect(f.metadata.replayTrimmedToTokens).toBe(42_000); + expect(f.metadata.replayOverflow).toBe(true); + }); + + it('omits the replay metadata when not provided', () => { + const f = flushAssistant([], '', 'completed', { finishReason: 'stop' }); + expect('replayTrimmedToTokens' in f.metadata).toBe(false); + expect('replayOverflow' in f.metadata).toBe(false); + }); + // #274 observability: the page-change diff the agent saw this turn is persisted // to metadata.pageChanged when a non-empty diff was injected, and omitted when // the diff is empty/whitespace or the arg is not supplied. 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 4f8cb2dc..b4c65333 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -55,6 +55,12 @@ import { type SelectionContext, } from './tools/current-page.util'; import { roleModelOverride } from './roles/role-model-config'; +import { + resolveReplayBudget, + isContextOverflowError, + trimHistoryForReplay, + REPLAY_AGGRESSIVE_FRACTION, +} from './history-budget'; import { startSseHeartbeat, stripStreamingHopByHopHeaders, @@ -127,6 +133,15 @@ const STEP_LIMIT_NO_ANSWER_MARKER = const OUTPUT_DEGENERATION_ERROR = 'Output degeneration detected (repeated token loop)'; +// Prefix recorded on the assistant row when the provider rejected the turn for +// CONTEXT OVERFLOW (#490): the replayed history exceeded the model's window. The +// row is ALSO stamped `metadata.replayOverflow` so the NEXT turn's budgeter trims +// aggressively (the reactive recovery — the overflowing turn had no usage signal +// to trigger preventive trimming, so the classified 400 is what un-bricks it). +export const CONTEXT_OVERFLOW_ERROR_PREFIX = + 'Диалог превысил контекстное окно модели; история будет агрессивно ' + + 'сокращена на следующем ходу.'; + /** * Compute the step-budget warning text (#444), or '' when this step is outside * the warning band. The warning fires on steps @@ -1086,7 +1101,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // per-row conversion and degraded to plain text with a "[tool context // omitted]" marker rather than 500-ing the whole turn (silent loss of tool // context is not acceptable — the model must see the truncation). - const messages = await convertHistoryResilient(uiMessages, (index, err) => + let messages = await convertHistoryResilient(uiMessages, (index, err) => this.logger.warn( `Degraded unconvertible history row ${index} on chat ${chatId} to text: ${ err instanceof Error ? err.message : 'unknown error' @@ -1135,6 +1150,58 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // Here we only need the admin-configured system prompt. const resolved = await this.aiSettings.resolve(workspace.id); + // History-replay token budget (#490). The full conversation is replayed to + // the provider every turn, so a long chat eventually 400s on the context + // window — forever. Bound the REPLAYED history (never the persisted rows). + // PRIMARY signal is the provider's own fact: the last turn's contextTokens. + const replayBudget = resolveReplayBudget(resolved?.chatContextWindowRaw); + if (replayBudget.usedDefault) { + // The default fires precisely for installs with NO configured window — + // the ones that hit terminal overflow. Warn so it is observable. + this.logger.warn( + `AI chat (chat ${chatId}): no chatContextWindow configured; ` + + `applying the default replay budget (${replayBudget.thresholdTokens} tokens).`, + ); + } + // Last turn's provider-reported context size (authoritative when present). + const priorContextTokens = lastAssistantContextTokens(oldHistory); + // Reactive recovery (#490): if the LAST turn was rejected for context + // overflow (stamped by onError), trim AGGRESSIVELY this turn — the + // 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; + if (priorOverflowed) { + this.logger.warn( + `AI chat (chat ${chatId}): previous turn hit context overflow; ` + + `applying aggressive replay budget (${effectiveThreshold} tokens).`, + ); + } + const preTrim = trimHistoryForReplay( + messages, + effectiveThreshold, + // A prior OVERFLOW means the provider count is stale/absent — force the + // char-estimate path by ignoring priorContextTokens on recovery. + priorOverflowed ? undefined : priorContextTokens, + ); + messages = preTrim.messages; + // Observability (#490): record the budgeter's decision on the turn so the UI + // can surface "replay truncated at N tokens". Threaded into flushAssistant. + let replayTrimmedToTokens: number | undefined = preTrim.trimmed + ? preTrim.estimatedTokens + : undefined; + if (preTrim.trimmed) { + this.logger.log( + `AI chat (chat ${chatId}): replay history trimmed to ~${preTrim.estimatedTokens} ` + + `tokens (budget ${replayBudget.thresholdTokens}).`, + ); + } + // Build the external MCP toolset FIRST so the system prompt can carry each // connected server's admin-authored guidance (#180). Merge in admin- // configured external MCP tools (web search, etc.; §6.8). A down/slow @@ -1658,6 +1725,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { maxContextTokens: resolved?.chatContextWindow, pageChanged, partsCache, + replayTrimmedToTokens, }), ); // #184/#487: the RUN is finalized ALWAYS (never gated on the message). @@ -1705,7 +1773,16 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // object, so the actual provider cause is clearly logged. Reuse the // shared formatter so provider error formatting stays unified. const e = error as { stack?: string }; - const errorText = describeProviderError(error, String(error)); + // #490 reactive branch: classify a CONTEXT-OVERFLOW rejection (the + // replayed history exceeded the model window). The overflowing turn had + // no prior usage to trigger preventive trimming, so we record a clear, + // distinguishable cause AND stamp the row so the NEXT turn's budgeter + // trims aggressively — the reactive recovery that un-bricks the chat. + const overflow = isContextOverflowError(error); + const providerError = describeProviderError(error, String(error)); + const errorText = overflow + ? `${CONTEXT_OVERFLOW_ERROR_PREFIX} (${providerError})` + : providerError; this.logger.error(`AI chat stream error: ${errorText}`, e?.stack); // DIAGNOSTIC (Safari stream-drop investigation) — temporary: timing of // an error-terminated stream. @@ -1724,6 +1801,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { error: errorText, pageChanged, partsCache, + replayTrimmedToTokens, + replayOverflow: overflow || undefined, }), ); // #184: settle the RUN as failed, carrying the provider/transport cause. @@ -2103,6 +2182,45 @@ export function chatStreamMetadata( return undefined; } +/** + * The provider-reported context size of the most recent assistant turn, read from + * its persisted `metadata.contextTokens` (#490 replay budgeter's PRIMARY signal — + * the provider's own fact, not an estimate). Returns undefined for a chat with no + * assistant turn yet, or one whose last turn recorded no usage (e.g. it errored), + * in which case the budgeter falls back to the char-estimate. + */ +export function lastAssistantContextTokens( + history: ReadonlyArray, +): number | undefined { + for (let i = history.length - 1; i >= 0; i--) { + const row = history[i]; + if (row.role !== 'assistant') continue; + const meta = (row.metadata ?? {}) as { contextTokens?: unknown }; + const n = meta.contextTokens; + return typeof n === 'number' && Number.isFinite(n) && n > 0 ? n : undefined; + } + return undefined; +} + +/** + * Whether the most recent assistant turn was rejected for CONTEXT OVERFLOW + * (#490): its row carries `metadata.replayOverflow` (stamped by the stream's + * onError). The next turn's budgeter reads this to trim aggressively — the + * reactive recovery. Only the LAST assistant turn matters (an older overflow was + * already recovered), so we stop at the first assistant row scanning backwards. + */ +export function lastAssistantReplayOverflow( + history: ReadonlyArray, +): boolean { + for (let i = history.length - 1; i >= 0; i--) { + const row = history[i]; + if (row.role !== 'assistant') continue; + const meta = (row.metadata ?? {}) as { replayOverflow?: unknown }; + return meta.replayOverflow === true; + } + return false; +} + /** The last message with role 'user' from a useChat payload, if any. */ function lastUserMessage( messages: UIMessage[] | undefined, @@ -2636,6 +2754,13 @@ export function flushAssistant( // Per-turn step->parts memo (#490): pass the SAME cache on every flush of a // turn so each finished step's output is stringified once, not once per flush. partsCache?: StepPartsCache; + // #490 observability: when the replay budgeter trimmed this turn's history, + // the (estimated) token size it trimmed to — the UI can show "replay truncated + // at N tokens". Omitted when nothing was trimmed. + replayTrimmedToTokens?: number; + // #490 reactive branch: set when the provider rejected this turn for context + // overflow. Stamped into metadata so the NEXT turn's budgeter trims aggressively. + replayOverflow?: boolean; }, ): AssistantFlush { const finished = capturedSteps ?? []; @@ -2674,6 +2799,9 @@ export function flushAssistant( if (extra?.contextTokens) metadata.contextTokens = extra.contextTokens; if (extra?.maxContextTokens) metadata.maxContextTokens = extra.maxContextTokens; + if (extra?.replayTrimmedToTokens) + metadata.replayTrimmedToTokens = extra.replayTrimmedToTokens; + if (extra?.replayOverflow) metadata.replayOverflow = true; if (extra?.error) metadata.error = extra.error; // Persist the page-change diff the agent saw this turn (#274 observability), // so history / the Markdown export can show what the user changed. Only when 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.spec.ts new file mode 100644 index 00000000..cbd12ad2 --- /dev/null +++ b/apps/server/src/core/ai-chat/ai-chat.write-volume.spec.ts @@ -0,0 +1,209 @@ +import { randomBytes } from 'crypto'; +import { Client } from 'pg'; +import { flushAssistant, serializeSteps } from './ai-chat.service'; + +/** + * #490 write-volume regression — an OBSERVABLE-PROPERTY test on a LIVE Postgres, + * not "bytes through a mock repo" (a mock measures exactly the thing that does not + * hurt). It drives a realistic 50-step run where each step returns a ~100 KB tool + * output and, at every `onStepFinish`, UPDATEs the assistant row the way the + * service does — then reads the REAL write volume via the `pg_current_wal_lsn()` + * delta around the run. + * + * The property proven: v2 stores each tool OUTPUT only in `metadata.parts`, no + * longer ALSO in the `tool_calls` trace. So: + * 1. the trace (`tool_calls`) column's write volume is now O(Σ steps) — tiny, + * linear outcome flags — vs the pre-#490 O(N²) that re-persisted every prior + * output on every step; and + * 2. the FULL-row write volume drops sharply (the duplicated output copy is gone). + * + * Connects to the local gitmost test Postgres (docker `gitmost-test-pg` on :5432); + * SKIPS cleanly when that DB is not reachable so it never breaks a DB-less CI. + */ +const CONN = + process.env.WAL_TEST_DATABASE_URL ?? + 'postgresql://docmost:docmost_dev_pw@localhost:5432/docmost'; + +// A step whose tool output is ~100 KB (a page read), in the SDK StepLike shape. +// The body is INCOMPRESSIBLE random text — a `'x'.repeat()` filler would TOAST- +// compress to nothing and hide the real write volume (a page body does not). +function makeStep(i: number, outputBytes = 100_000) { + const body = randomBytes(Math.ceil(outputBytes * 0.75)).toString('base64'); + return { + text: `step ${i} reasoning`, + toolCalls: [{ toolCallId: `c${i}`, toolName: 'getPage', input: { id: `p${i}` } }], + toolResults: [ + { + toolCallId: `c${i}`, + toolName: 'getPage', + output: { id: `p${i}`, title: `Page ${i}`, body }, + }, + ], + }; +} + +// The pre-#490 (v1) trace: outputs stored a SECOND time in `tool_calls` +// (the duplication #490 removed). Mirrors the OLD serializeSteps shape. +function v1Trace(steps: ReturnType[]): unknown { + const calls: unknown[] = []; + for (const s of steps) { + for (const c of s.toolCalls) calls.push({ toolName: c.toolName, input: c.input }); + for (const r of s.toolResults) + calls.push({ toolName: r.toolName, output: r.output }); + } + return calls; +} + +async function walDelta( + client: Client, + fn: () => Promise, +): Promise { + const before = (await client.query('SELECT pg_current_wal_lsn() AS l')).rows[0] + .l as string; + await fn(); + // NOTE: do NOT pg_switch_wal() here — a segment switch pads the LSN to the next + // 16 MB boundary and would swamp the actual write delta. The raw LSN advances by + // the bytes of WAL emitted, which is exactly what we want to measure. + const after = (await client.query('SELECT pg_current_wal_lsn() AS l')).rows[0] + .l as string; + return Number( + (await client.query('SELECT pg_wal_lsn_diff($1,$2) AS d', [after, before])) + .rows[0].d, + ); +} + +describe('#490 write-volume on a live Postgres (pg_current_wal_lsn delta)', () => { + let client: Client | undefined; + let available = false; + + beforeAll(async () => { + try { + client = new Client(CONN); + await client.connect(); + await client.query('SELECT pg_current_wal_lsn()'); + available = true; + } catch { + available = false; + client = undefined; + } + }); + + afterAll(async () => { + await client?.end().catch(() => undefined); + }); + + const STEPS = 50; + + it('v2 trace write volume is O(Σ steps) — a tiny fraction of the v1 duplicate', async () => { + if (!available || !client) { + console.warn('SKIP: gitmost-test-pg not reachable; skipping WAL test.'); + return; + } + const c = client; + // Isolated table so we measure only the tool_calls (trace) column's writes. + await c.query('DROP TABLE IF EXISTS _wal_trace'); + await c.query('CREATE TABLE _wal_trace(id int primary key, tool_calls jsonb)'); + await c.query("INSERT INTO _wal_trace VALUES (1, '[]'::jsonb)"); + + const steps: ReturnType[] = []; + + // v1: each step re-persists ALL prior outputs into the trace (the O(N²) churn). + const v1 = await walDelta(c, async () => { + const acc: ReturnType[] = []; + for (let i = 0; i < STEPS; i++) { + acc.push(makeStep(i)); + await c.query('UPDATE _wal_trace SET tool_calls=$1 WHERE id=1', [ + JSON.stringify(v1Trace(acc)), + ]); + } + steps.push(...acc); + }); + + await c.query("UPDATE _wal_trace SET tool_calls='[]'::jsonb WHERE id=1"); + + // v2: the REAL serializeSteps — outcome flags only, NO outputs. + const v2 = await walDelta(c, async () => { + const acc: ReturnType[] = []; + for (let i = 0; i < STEPS; i++) { + acc.push(makeStep(i)); + await c.query('UPDATE _wal_trace SET tool_calls=$1 WHERE id=1', [ + JSON.stringify(serializeSteps(acc)), + ]); + } + }); + + await c.query('DROP TABLE IF EXISTS _wal_trace'); + + // eslint-disable-next-line no-console + console.log( + `[#490 WAL] trace column over ${STEPS} steps: v1=${(v1 / 1e6).toFixed(1)}MB ` + + `v2=${(v2 / 1e6).toFixed(2)}MB (${(v1 / v2).toFixed(0)}x smaller)`, + ); + + // The trace no longer carries outputs: v2 is a tiny fraction of v1's WAL. + expect(v2).toBeLessThan(v1 * 0.1); + // And v2's trace WAL is small in absolute terms — O(Σ steps) of flags, not + // O(N² × output). 50 steps of ~40-byte flags is well under a few MB of WAL. + expect(v2).toBeLessThan(5_000_000); + // v1's duplicate alone is huge (≈ the 100 KB output re-written N² times). + expect(v1).toBeGreaterThan(50_000_000); + }, 120_000); + + it('the full assistant row write drops sharply once the duplicate is gone', async () => { + if (!available || !client) return; + const c = client; + await c.query('DROP TABLE IF EXISTS _wal_full'); + await c.query( + 'CREATE TABLE _wal_full(id int primary key, content text, tool_calls jsonb, metadata jsonb, status text)', + ); + await c.query("INSERT INTO _wal_full VALUES (1, '', '[]'::jsonb, '{}'::jsonb, 'streaming')"); + + const writeRow = async (patch: { + content: string; + toolCalls: unknown; + metadata: unknown; + status: string; + }) => + c.query( + 'UPDATE _wal_full SET content=$1, tool_calls=$2, metadata=$3, status=$4 WHERE id=1', + [ + patch.content, + JSON.stringify(patch.toolCalls ?? null), + JSON.stringify(patch.metadata), + patch.status, + ], + ); + + // v2 (real flushAssistant): outputs live once, in metadata.parts. + const v2 = await walDelta(c, async () => { + const acc: ReturnType[] = []; + for (let i = 0; i < STEPS; i++) { + acc.push(makeStep(i)); + await writeRow(flushAssistant(acc as never, '', 'streaming')); + } + }); + + await c.query("UPDATE _wal_full SET content='', tool_calls='[]'::jsonb, metadata='{}'::jsonb WHERE id=1"); + + // v1: same row PLUS the duplicated outputs in the trace column. + const v1 = await walDelta(c, async () => { + const acc: ReturnType[] = []; + for (let i = 0; i < STEPS; i++) { + acc.push(makeStep(i)); + const f = flushAssistant(acc as never, '', 'streaming'); + await writeRow({ ...f, toolCalls: v1Trace(acc) }); + } + }); + + await c.query('DROP TABLE IF EXISTS _wal_full'); + + // eslint-disable-next-line no-console + console.log( + `[#490 WAL] full row over ${STEPS} steps: v1=${(v1 / 1e6).toFixed(1)}MB ` + + `v2=${(v2 / 1e6).toFixed(1)}MB (saved ${((1 - v2 / v1) * 100).toFixed(0)}%)`, + ); + + // Removing the duplicated trace copy is a large, real write-volume reduction. + expect(v2).toBeLessThan(v1 * 0.75); + }, 120_000); +}); diff --git a/apps/server/src/core/ai-chat/history-budget.spec.ts b/apps/server/src/core/ai-chat/history-budget.spec.ts new file mode 100644 index 00000000..16f6fd76 --- /dev/null +++ b/apps/server/src/core/ai-chat/history-budget.spec.ts @@ -0,0 +1,266 @@ +import type { ModelMessage } from 'ai'; +import { + resolveReplayBudget, + isContextOverflowError, + estimateMessagesTokens, + trimHistoryForReplay, + REPLAY_BUDGET_DEFAULT_TOKENS, + REPLAY_TRUNCATION_MARKER, + REPLAY_TURN_COLLAPSED_MARKER, +} from './history-budget'; + +describe('resolveReplayBudget', () => { + it('uses min(default, 0.7 x window) for a configured window', () => { + // 0.7 x 60k = 42k < 100k + expect(resolveReplayBudget(60_000)).toEqual({ + thresholdTokens: 42_000, + usedDefault: false, + }); + // 0.7 x 1M = 700k, capped to the 100k default + expect(resolveReplayBudget(1_000_000)).toEqual({ + thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS, + usedDefault: false, + }); + }); + + it('accepts the raw ::text stored form', () => { + expect(resolveReplayBudget('60000').thresholdTokens).toBe(42_000); + }); + + // The crux (#490): a chat with NO context window configured must STILL be + // budgeted — those are exactly the installs that hit terminal overflow. + it('applies the flat default when the window is unset/empty', () => { + expect(resolveReplayBudget(undefined)).toEqual({ + thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS, + usedDefault: true, + }); + expect(resolveReplayBudget('')).toEqual({ + thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS, + usedDefault: true, + }); + expect(resolveReplayBudget(' ')).toEqual({ + thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS, + usedDefault: true, + }); + }); + + it('treats an explicit 0 as the off-switch (distinct from unset)', () => { + expect(resolveReplayBudget(0)).toEqual({ + thresholdTokens: null, + usedDefault: false, + }); + expect(resolveReplayBudget('0')).toEqual({ + thresholdTokens: null, + usedDefault: false, + }); + }); + + it('falls back to the default on a negative/garbage value', () => { + expect(resolveReplayBudget(-5).usedDefault).toBe(true); + expect(resolveReplayBudget('abc').usedDefault).toBe(true); + }); +}); + +describe('isContextOverflowError', () => { + it('classifies a real provider 400 context-overflow shape', () => { + // OpenAI-compatible shape. + expect( + isContextOverflowError({ + statusCode: 400, + message: + "This model's maximum context length is 128000 tokens. However, your messages resulted in 214000 tokens. Please reduce the length of the messages.", + }), + ).toBe(true); + // Anthropic-style wording. + expect( + isContextOverflowError({ + status: 400, + message: 'prompt is too long: 250000 tokens > 200000 maximum', + }), + ).toBe(true); + // Nested body + string status. + expect( + isContextOverflowError({ + response: { status: '400' }, + message: 'input is too long for the requested model', + }), + ).toBe(true); + // Error instance with the cause carrying the body. + const e = new Error('Bad request'); + (e as any).statusCode = 400; + (e as any).cause = new Error('maximum context window exceeded'); + expect(isContextOverflowError(e)).toBe(true); + }); + + it('does NOT classify unrelated 400s or auth/rate-limit errors', () => { + expect( + isContextOverflowError({ statusCode: 400, message: 'invalid tool schema' }), + ).toBe(false); + expect( + isContextOverflowError({ + statusCode: 429, + message: 'context length exceeded but rate limited', + }), + ).toBe(false); + expect(isContextOverflowError({ statusCode: 500, message: 'server error' })).toBe( + false, + ); + expect(isContextOverflowError(undefined)).toBe(false); + expect(isContextOverflowError('some random string')).toBe(false); + }); +}); + +// Helpers to build ModelMessage fixtures in the ai@6 shape. +const userMsg = (text: string): ModelMessage => + ({ role: 'user', content: [{ type: 'text', text }] }) as ModelMessage; +const assistantMsg = ( + text: string, + toolCallId?: string, + toolName?: string, +): ModelMessage => + ({ + role: 'assistant', + content: [ + { type: 'text', text }, + ...(toolCallId + ? [{ type: 'tool-call', toolCallId, toolName, input: {} }] + : []), + ], + }) as ModelMessage; +const toolMsg = ( + toolCallId: string, + toolName: string, + value: unknown, +): ModelMessage => + ({ + role: 'tool', + content: [ + { type: 'tool-result', toolCallId, toolName, output: { type: 'json', value } }, + ], + }) as ModelMessage; + +describe('trimHistoryForReplay', () => { + it('null budget disables trimming (returns the same reference)', () => { + const msgs = [userMsg('hi'), assistantMsg('yo')]; + const r = trimHistoryForReplay(msgs, null); + expect(r.trimmed).toBe(false); + expect(r.messages).toBe(msgs); + }); + + it('leaves history under budget untouched (same reference)', () => { + const msgs = [userMsg('hi'), assistantMsg('a short answer')]; + const r = trimHistoryForReplay(msgs, 100_000); + expect(r.trimmed).toBe(false); + expect(r.messages).toBe(msgs); + }); + + it('truncates OLD tool outputs but keeps recent turns full', () => { + const big = 'X'.repeat(40_000); // ~16k tokens on its own + const msgs: ModelMessage[] = []; + // 6 OLD turns (indices 0..5), each with a huge tool output. + for (let i = 0; i < 6; i++) { + msgs.push(userMsg(`old q${i}`)); + msgs.push(assistantMsg('looking', `c${i}`, 'getPage')); + msgs.push(toolMsg(`c${i}`, 'getPage', { body: big })); + msgs.push(assistantMsg(`old a${i}`)); + } + // 3 small recent turns, then the CURRENT turn with its own huge tool output. + // With REPLAY_KEEP_RECENT_TURNS=4 the last 4 user-turns stay full, so only + // these small recent turns + the current big one are kept full; the 6 old + // turns above fall in the trim region. + for (let i = 0; i < 3; i++) { + msgs.push(userMsg(`recent q${i}`)); + msgs.push(assistantMsg(`recent a${i}`)); + } + msgs.push(userMsg('current q')); + msgs.push(assistantMsg('looking', 'cR', 'getPage')); + msgs.push(toolMsg('cR', 'getPage', { body: big })); + msgs.push(assistantMsg('current a')); + + // Budget large enough that phase-1 tool truncation alone brings it under. + const r = trimHistoryForReplay(msgs, 30_000); + expect(r.trimmed).toBe(true); + const flat = JSON.stringify(r.messages); + // The CURRENT turn's tool output survives in full. + expect(flat).toContain(big); + // Old outputs were truncated with the marker. + expect(flat).toContain(REPLAY_TRUNCATION_MARKER); + // Phase 1 sufficed: the oldest turns were NOT collapsed. + expect(flat).not.toContain(REPLAY_TURN_COLLAPSED_MARKER); + expect(estimateMessagesTokens(r.messages)).toBeLessThan( + estimateMessagesTokens(msgs), + ); + }); + + it('collapses the oldest turns when tool truncation is not enough', () => { + // Many turns with LARGE assistant TEXT (not tool output) so phase 1 can't help. + const bigText = 'слово '.repeat(8_000); // large Cyrillic text per turn + const msgs: ModelMessage[] = []; + for (let i = 0; i < 12; i++) { + msgs.push(userMsg(`q${i}`)); + msgs.push(assistantMsg(bigText)); + } + const r = trimHistoryForReplay(msgs, 30_000); + expect(r.trimmed).toBe(true); + // Oldest turns collapsed; result fits (best-effort) and is much smaller. + expect(estimateMessagesTokens(r.messages)).toBeLessThan( + estimateMessagesTokens(msgs), + ); + // The LAST turn's text is preserved in full (recent turns stay full). + expect(JSON.stringify(r.messages[r.messages.length - 1])).toContain(bigText); + }); + + it('is deterministic / byte-stable for identical inputs', () => { + const big = 'Y'.repeat(30_000); + const build = (): ModelMessage[] => { + const m: ModelMessage[] = []; + for (let i = 0; i < 10; i++) { + m.push(userMsg(`q${i}`)); + m.push(assistantMsg('t', `c${i}`, 'getPage')); + m.push(toolMsg(`c${i}`, 'getPage', { body: big })); + } + return m; + }; + const a = trimHistoryForReplay(build(), 15_000); + const b = trimHistoryForReplay(build(), 15_000); + expect(JSON.stringify(a.messages)).toBe(JSON.stringify(b.messages)); + }); + + it('never leaves an unpaired tool-call after collapsing (balanced history)', () => { + const big = 'Z'.repeat(40_000); + const msgs: ModelMessage[] = []; + for (let i = 0; i < 10; i++) { + msgs.push(userMsg(`q${i}`)); + msgs.push(assistantMsg('t', `c${i}`, 'getPage')); + msgs.push(toolMsg(`c${i}`, 'getPage', { body: big })); + } + const r = trimHistoryForReplay(msgs, 8_000); + // Count tool-call vs tool-result parts in the trimmed output. + let calls = 0; + let results = 0; + for (const m of r.messages) { + if (!Array.isArray(m.content)) continue; + for (const p of m.content as Array<{ type?: string }>) { + if (p.type === 'tool-call') calls++; + if (p.type === 'tool-result' || p.type === 'tool-error') results++; + } + } + // Every surviving tool-call has a surviving result (collapsing drops BOTH). + expect(calls).toBe(results); + // Collapsed turns carry the marker. + expect(JSON.stringify(r.messages)).toContain(REPLAY_TURN_COLLAPSED_MARKER); + }); + + it('respects the provider fact: under-budget contextTokens skips trimming', () => { + const big = 'W'.repeat(60_000); + const msgs = [ + userMsg('q'), + assistantMsg('t', 'c1', 'getPage'), + toolMsg('c1', 'getPage', { body: big }), + ]; + // char-estimate is high, but the provider says we are well under budget. + const r = trimHistoryForReplay(msgs, 100_000, 5_000); + expect(r.trimmed).toBe(false); + expect(r.messages).toBe(msgs); + }); +}); diff --git a/apps/server/src/core/ai-chat/history-budget.ts b/apps/server/src/core/ai-chat/history-budget.ts new file mode 100644 index 00000000..fb45702f --- /dev/null +++ b/apps/server/src/core/ai-chat/history-budget.ts @@ -0,0 +1,356 @@ +/** + * History-replay token budget (#490). + * + * The whole persisted conversation is replayed to the provider on EVERY turn, so + * a long chat eventually exceeds the model's context window and the provider 400s + * on every turn — terminally (the chat "bricks"). This module bounds the replayed + * history at REPLAY TIME only: it never mutates what is persisted (the DB stays + * the full record), and its output is a deterministic, byte-stable function of its + * input so the trimmed prefix is identical turn to turn (provider prompt-cache + * friendliness — real money on long chats). + * + * The PRIMARY signal is the provider's own fact: `metadata.contextTokens` from the + * last turn. The chars-based {@link estimateTokens} (shared with the client) is + * used only for the DELTA of not-yet-sent messages, to decide WHAT to trim, and as + * the fallback for chats with no usage yet. + */ +import type { ModelMessage } from 'ai'; +import { estimateTokens } from '@docmost/token-estimate'; + +/** Flat default budget when no context window is configured (tokens). */ +export const REPLAY_BUDGET_DEFAULT_TOKENS = 100_000; +/** Fraction of a configured context window used as the budget. */ +export const REPLAY_BUDGET_WINDOW_FRACTION = 0.7; +/** + * Fraction of the normal budget used for the REACTIVE re-trim after a provider + * context-overflow 400 — the preventive estimate under-counted, so cut harder. + */ +export const REPLAY_AGGRESSIVE_FRACTION = 0.5; +/** + * Turns (a user message + its assistant/tool replies) kept FULL at the tail, + * including the current one — never trimmed. Older turns are compacted first. + */ +export const REPLAY_KEEP_RECENT_TURNS = 4; +/** Leading chars kept from a truncated old tool output. */ +export const REPLAY_TOOL_OUTPUT_HEAD = 800; +/** Trailing chars kept from a truncated old tool output. */ +export const REPLAY_TOOL_OUTPUT_TAIL = 300; +/** Marker inserted where an old tool output was truncated for replay. */ +export const REPLAY_TRUNCATION_MARKER = + '[…truncated for replay; call the tool again to read the full output]'; +/** Marker for a whole old turn collapsed to its text. */ +export const REPLAY_TURN_COLLAPSED_MARKER = + '[earlier tool activity omitted for replay]'; + +export interface ReplayBudget { + /** Token threshold above which replay history is trimmed; `null` = OFF. */ + thresholdTokens: number | null; + /** True when the flat default was used (no context window configured). */ + usedDefault: boolean; +} + +/** + * Resolve the replay budget from the RAW stored `chatContextWindow` (text/number). + * - a positive value -> `min(default, floor(fraction × window))` + * - 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) + * + * Note the raw value is needed because the parsed `chatContextWindow` collapses + * both `0` and unset to `undefined`, which would erase the explicit off-switch. + */ +export function resolveReplayBudget(rawContextWindow: unknown): ReplayBudget { + let n: number | undefined; + if (typeof rawContextWindow === 'number') { + n = rawContextWindow; + } else if (typeof rawContextWindow === 'string') { + const t = rawContextWindow.trim(); + n = t === '' ? undefined : Number(t); + } + // Unset / empty / non-numeric / negative -> flat default (the protective case). + if (n === undefined || !Number.isFinite(n) || n < 0) { + return { thresholdTokens: REPLAY_BUDGET_DEFAULT_TOKENS, usedDefault: true }; + } + // Explicit 0 -> off-switch. + if (n === 0) { + return { thresholdTokens: null, usedDefault: false }; + } + return { + thresholdTokens: Math.min( + REPLAY_BUDGET_DEFAULT_TOKENS, + Math.floor(REPLAY_BUDGET_WINDOW_FRACTION * n), + ), + usedDefault: false, + }; +} + +/** + * 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 + * message; match both the status and the message patterns robustly across + * OpenAI-compatible / Anthropic / Gemini wordings, since the exact shape varies. + */ +export function isContextOverflowError(error: unknown): boolean { + const status = extractStatus(error); + const msg = extractMessage(error).toLowerCase(); + // Message patterns seen across providers for "prompt too long". + const overflowPattern = + /context (?:length|window)|maximum context|too many tokens|too large for|reduce the length|prompt is too long|input (?:is )?too long|exceeds? the (?:maximum )?(?:context|token)|maximum.*tokens|string too long/; + if (!overflowPattern.test(msg)) return false; + // A 400/413 with an overflow-shaped message is an overflow. Some providers + // omit/rewrite the status, so accept the message match when the status is + // unknown, but reject it for auth/rate-limit statuses that never mean overflow. + if (status === 400 || status === 413) return true; + if (status === 401 || status === 403 || status === 429) return false; + return true; +} + +function extractStatus(error: unknown): number | undefined { + if (!error || typeof error !== 'object') return undefined; + const e = error as Record; + for (const k of ['statusCode', 'status']) { + const v = e[k]; + if (typeof v === 'number') return v; + if (typeof v === 'string' && /^\d+$/.test(v)) return Number(v); + } + // Nested (e.g. { response: { status } } / { cause: { statusCode } }). + for (const k of ['response', 'cause', 'data']) { + const nested = e[k]; + if (nested && typeof nested === 'object') { + const s = extractStatus(nested); + if (s !== undefined) return s; + } + } + return undefined; +} + +function extractMessage(error: unknown): string { + if (error == null) return ''; + if (typeof error === 'string') return error; + if (error instanceof Error) { + // Include nested causes (provider libs wrap the real body in `cause`). + const cause = (error as { cause?: unknown }).cause; + return `${error.message} ${cause ? extractMessage(cause) : ''}`; + } + if (typeof error === 'object') { + const e = error as Record; + const parts: string[] = []; + for (const k of ['message', 'error', 'body', 'responseBody', 'data']) { + const v = e[k]; + if (typeof v === 'string') parts.push(v); + else if (v && typeof v === 'object') parts.push(extractMessage(v)); + } + return parts.join(' '); + } + return String(error); +} + +/** Rough token size of a ModelMessage array via the shared chars estimator. */ +export function estimateMessagesTokens( + messages: ReadonlyArray, +): number { + let total = 0; + for (const m of messages) { + total += estimateTokens(serializeContent(m.content)); + } + return total; +} + +function serializeContent(content: unknown): string { + if (typeof content === 'string') return content; + try { + return JSON.stringify(content) ?? ''; + } catch { + return ''; + } +} + +/** Deep JSON string of an arbitrary value, bounded so estimation never throws. */ +function stringifyValue(value: unknown): string { + if (typeof value === 'string') return value; + try { + return JSON.stringify(value) ?? String(value); + } catch { + return String(value); + } +} + +export interface TrimResult { + messages: ModelMessage[]; + /** Whether any trimming was applied. */ + trimmed: boolean; + /** Estimated tokens of the returned messages (chars-based). */ + estimatedTokens: number; +} + +/** + * Bound the replayed history to `budgetTokens`, deterministically. Returns the + * SAME array reference (no copy) when nothing needs trimming, so the common case + * is free and byte-identical. Trimming order (spec #490): + * 1. truncate OLD turns' tool outputs (head+tail + marker) — the bulk of the size + * 2. mechanically collapse the OLDEST turns to their text (concatenation, no LLM) + * 3. the current + last {@link REPLAY_KEEP_RECENT_TURNS} turns stay FULL + * + * `budgetTokens === null` disables trimming. `priorContextTokens` (the provider's + * fact from last turn) short-circuits the decision: when it is known and already + * under budget we skip trimming even if the char-estimate is higher (the provider + * count is authoritative). The char-estimate drives WHAT to cut. + */ +export function trimHistoryForReplay( + messages: ModelMessage[], + budgetTokens: number | null, + priorContextTokens?: number, +): TrimResult { + if (budgetTokens == null) { + return { messages, trimmed: false, estimatedTokens: 0 }; + } + const estimated = estimateMessagesTokens(messages); + // Decision signal: prefer the provider's fact (last turn's contextTokens) plus + // the estimated delta of the messages appended since; fall back to the pure + // char-estimate for a chat with no usage yet. + const projected = + priorContextTokens != null + ? Math.max(priorContextTokens, estimated) + : estimated; + if (projected <= budgetTokens) { + return { messages, trimmed: false, estimatedTokens: estimated }; + } + + // The tail we always keep full: from the Nth-from-last user message onward. + const boundary = recentBoundaryIndex(messages, REPLAY_KEEP_RECENT_TURNS); + const tail = messages.slice(boundary); + let head = messages.slice(0, boundary).map(cloneMessage); + + // Phase 1: truncate old tool outputs. + for (const m of head) { + if (m.role === 'tool') truncateToolMessage(m); + } + let out = [...head, ...tail]; + let est = estimateMessagesTokens(out); + if (est <= budgetTokens) { + return { messages: out, trimmed: true, estimatedTokens: est }; + } + + // Phase 2: collapse the oldest turns (in `head`) to their text, one at a time, + // from the oldest, until we fit or the whole head is collapsed. + const turns = splitTurns(head); + const collapsed: ModelMessage[] = []; + let i = 0; + for (; i < turns.length; i++) { + if (est <= budgetTokens) break; + collapsed.push(...collapseTurn(turns[i])); + // Re-estimate the whole prospective output. + const remaining = turns.slice(i + 1).flat(); + out = [...collapsed, ...remaining, ...tail]; + est = estimateMessagesTokens(out); + } + // Include any turns we didn't need to collapse. + const remaining = turns.slice(i).flat(); + out = [...collapsed, ...remaining, ...tail]; + est = estimateMessagesTokens(out); + return { messages: out, trimmed: true, estimatedTokens: est }; +} + +/** Index of the first message of the Nth-from-last user turn (0 if fewer). */ +function recentBoundaryIndex( + messages: ReadonlyArray, + keepTurns: number, +): number { + const userIdx: number[] = []; + for (let i = 0; i < messages.length; i++) { + if (messages[i].role === 'user') userIdx.push(i); + } + if (userIdx.length <= keepTurns) return 0; + return userIdx[userIdx.length - keepTurns]; +} + +/** Split a message list into turns; each turn starts at a `user` message. */ +function splitTurns(messages: ModelMessage[]): ModelMessage[][] { + const turns: ModelMessage[][] = []; + for (const m of messages) { + if (m.role === 'user' || turns.length === 0) turns.push([m]); + else turns[turns.length - 1].push(m); + } + return turns; +} + +/** + * Collapse a whole turn to its plain text (mechanical concatenation, not an LLM + * summary). Keeps the user message; replaces the assistant/tool messages with a + * single assistant text message = the assistant's concatenated text + a marker + * when tool activity was dropped. Dropping BOTH the tool-call and tool-result + * parts together keeps the rebuilt history balanced (no unpaired calls). + */ +function collapseTurn(turn: ModelMessage[]): ModelMessage[] { + const out: ModelMessage[] = []; + let assistantText = ''; + let hadTools = false; + for (const m of turn) { + if (m.role === 'user') { + out.push(m); + } else if (m.role === 'assistant') { + const { text, tools } = extractAssistantText(m.content); + assistantText += text; + hadTools = hadTools || tools; + } else if (m.role === 'tool') { + hadTools = true; + } else { + out.push(m); + } + } + const note = + (assistantText ? assistantText.trimEnd() : '') + + (hadTools + ? `${assistantText ? '\n\n' : ''}${REPLAY_TURN_COLLAPSED_MARKER}` + : ''); + if (note) out.push({ role: 'assistant', content: note } as ModelMessage); + return out; +} + +function extractAssistantText(content: unknown): { + text: string; + tools: boolean; +} { + if (typeof content === 'string') return { text: content, tools: false }; + if (!Array.isArray(content)) return { text: '', tools: false }; + let text = ''; + let tools = false; + for (const part of content) { + const type = (part as { type?: string })?.type; + if (type === 'text') text += (part as { text?: string }).text ?? ''; + else if (type === 'tool-call') tools = true; + } + return { text, tools }; +} + +/** Truncate every tool-result output in a `tool` message to head+tail+marker. */ +function truncateToolMessage(message: ModelMessage): void { + const content = message.content; + if (!Array.isArray(content)) return; + for (const part of content) { + const p = part as { type?: string; output?: { type?: string; value?: unknown } }; + if (p.type !== 'tool-result' && p.type !== 'tool-error') continue; + if (!p.output) continue; + const raw = stringifyValue(p.output.value); + const budget = REPLAY_TOOL_OUTPUT_HEAD + REPLAY_TOOL_OUTPUT_TAIL; + if (raw.length <= budget + REPLAY_TRUNCATION_MARKER.length) continue; + const truncated = + raw.slice(0, REPLAY_TOOL_OUTPUT_HEAD) + + `\n${REPLAY_TRUNCATION_MARKER}\n` + + raw.slice(raw.length - REPLAY_TOOL_OUTPUT_TAIL); + // Represent the shrunk output as a text output (a valid tool-result output). + p.output = { type: 'text', value: truncated }; + } +} + +/** Shallow-ish clone so trimming never mutates the caller's (persisted-derived) + * message objects — only the OLD region is cloned before it is edited. */ +function cloneMessage(m: ModelMessage): ModelMessage { + if (typeof m.content === 'string') return { ...m }; + return { + ...m, + content: (m.content as unknown[]).map((p) => + p && typeof p === 'object' ? { ...(p as object) } : p, + ), + } as ModelMessage; +} diff --git a/apps/server/src/integrations/ai/ai-settings.service.ts b/apps/server/src/integrations/ai/ai-settings.service.ts index 6be0b5be..ad98c256 100644 --- a/apps/server/src/integrations/ai/ai-settings.service.ts +++ b/apps/server/src/integrations/ai/ai-settings.service.ts @@ -245,6 +245,9 @@ export class AiSettingsService { // Max context window for the chat header badge denominator. Stored as // ::text; 0/unset/invalid = no limit (undefined). chatContextWindow: parsePositiveInt(provider.chatContextWindow), + // RAW stored value (#490): the replay budgeter reads this to distinguish an + // explicit `0` (off-switch) from unset, which parsePositiveInt cannot. + chatContextWindowRaw: provider.chatContextWindow, // Plain passthrough; getChatModel defaults unset to 'openai-compatible'. chatApiStyle: provider.chatApiStyle, // Cheap model id for the anonymous public-share assistant; reuses the chat diff --git a/apps/server/src/integrations/ai/ai.types.ts b/apps/server/src/integrations/ai/ai.types.ts index 06bf83e3..c43d2dfc 100644 --- a/apps/server/src/integrations/ai/ai.types.ts +++ b/apps/server/src/integrations/ai/ai.types.ts @@ -105,6 +105,10 @@ export interface ResolvedAiConfig extends Partial { // Max context window in tokens; surfaced to the chat header badge as the // "current / max" denominator. 0/unset = no limit. chatContextWindow?: number; + // RAW stored context window (::text), BEFORE parsePositiveInt collapses `0` and + // unset to `undefined`. The #490 replay budgeter needs the raw value to honor an + // explicit `0` off-switch distinctly from "unset -> flat default". + chatContextWindowRaw?: string | number; // Cheap model id for the public-share assistant; reuses the chat creds. publicShareChatModel?: string; // Agent-role id whose persona the public-share assistant adopts (empty/unset diff --git a/packages/token-estimate/package.json b/packages/token-estimate/package.json new file mode 100644 index 00000000..eef030a6 --- /dev/null +++ b/packages/token-estimate/package.json @@ -0,0 +1,19 @@ +{ + "name": "@docmost/token-estimate", + "version": "0.1.0", + "description": "Shared, provider-agnostic token estimator (chars/2.5) used by the AI-chat client counter and the server history-replay budgeter, so the two never diverge.", + "private": true, + "main": "./dist/index.js", + "types": "./dist/index.d.ts", + "scripts": { + "build": "tsc", + "watch": "tsc --watch", + "test": "vitest run", + "test:watch": "vitest" + }, + "license": "MIT", + "devDependencies": { + "typescript": "^5.0.0", + "vitest": "4.1.6" + } +} diff --git a/packages/token-estimate/src/index.test.ts b/packages/token-estimate/src/index.test.ts new file mode 100644 index 00000000..2854efd0 --- /dev/null +++ b/packages/token-estimate/src/index.test.ts @@ -0,0 +1,31 @@ +import { describe, it, expect } from 'vitest'; +import { estimateTokens, CHARS_PER_TOKEN } from './index'; + +describe('estimateTokens (shared chars/2.5)', () => { + it('returns 0 for empty / nullish input', () => { + expect(estimateTokens('')).toBe(0); + expect(estimateTokens(null)).toBe(0); + expect(estimateTokens(undefined)).toBe(0); + }); + + it('uses the chars/2.5 ratio, ceiled', () => { + expect(CHARS_PER_TOKEN).toBe(2.5); + // 5 chars / 2.5 = 2 + expect(estimateTokens('abcde')).toBe(2); + // any non-empty string is at least 1 token (ceil) + expect(estimateTokens('a')).toBe(1); + // 100 chars / 2.5 = 40 + expect(estimateTokens('x'.repeat(100))).toBe(40); + }); + + it('counts Cyrillic ~2x higher than the old chars/4 rule (no undercount)', () => { + const cyr = 'привет мир как дела'; // 19 chars + expect(estimateTokens(cyr)).toBe(Math.ceil(19 / 2.5)); // 8 + expect(estimateTokens(cyr)).toBeGreaterThan(Math.ceil(19 / 4)); // > 5 + }); + + it('is deterministic / byte-stable (same input => same output)', () => { + const s = 'the quick brown fox'; + expect(estimateTokens(s)).toBe(estimateTokens(s)); + }); +}); diff --git a/packages/token-estimate/src/index.ts b/packages/token-estimate/src/index.ts new file mode 100644 index 00000000..ec8d2c04 --- /dev/null +++ b/packages/token-estimate/src/index.ts @@ -0,0 +1,35 @@ +/** + * Shared, provider-agnostic token estimator (#490). + * + * No provider exposes an exact tokenizer we can afford to run on the hot path (a + * real BPE pass is O(n²)-ish, bloats the client bundle, and is wrong for + * Gemini/Ollama anyway), so both the client's in-body counter AND the server's + * history-replay budgeter use this ONE cheap chars-based heuristic. Keeping it in + * a single shared module is deliberate: two independent estimators drift, and then + * "the badge shows 60%" while "the budgeter already trimmed" — the exact confusion + * this package prevents. + * + * Ratio: **chars / 2.5**. Most content here is Cyrillic, where a token is ~2.5 + * characters; the common English `chars/4` rule of thumb UNDER-counts Cyrillic by + * ~2×, which for a budget check is the dangerous direction (it lets the context + * overflow). 2.5 slightly over-estimates pure English/code, which is the SAFE + * direction for a budget. This is an estimate, never an exact count — the + * authoritative figure is always the provider's reported usage; the estimate is + * for UI affordances, the delta of not-yet-sent messages, and deciding what to + * trim. + */ + +/** Characters per token for the shared estimate. See the module comment. */ +export const CHARS_PER_TOKEN = 2.5; + +/** + * Rough token estimate for a piece of text (chars / {@link CHARS_PER_TOKEN}). + * Returns 0 for empty/nullish input, and ceils so any non-empty text counts as at + * least one token. Pure and deterministic (byte-stable), so the same text always + * yields the same estimate — which the server budgeter relies on to keep replay + * trimming stable turn to turn (provider prompt-cache friendliness). + */ +export function estimateTokens(text: string | null | undefined): number { + if (!text) return 0; + return Math.ceil(text.length / CHARS_PER_TOKEN); +} diff --git a/packages/token-estimate/tsconfig.json b/packages/token-estimate/tsconfig.json new file mode 100644 index 00000000..dec9218d --- /dev/null +++ b/packages/token-estimate/tsconfig.json @@ -0,0 +1,16 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "CommonJS", + "moduleResolution": "Node", + "outDir": "./dist", + "rootDir": "./src", + "strict": true, + "esModuleInterop": true, + "skipLibCheck": true, + "forceConsistentCasingInFileNames": true, + "declaration": true + }, + "include": ["src/**/*"], + "exclude": ["src/**/*.test.ts"] +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index bdf58b2b..c876808e 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -287,6 +287,9 @@ importers: '@docmost/prosemirror-markdown': specifier: workspace:* version: link:../../packages/prosemirror-markdown + '@docmost/token-estimate': + specifier: workspace:* + version: link:../../packages/token-estimate '@excalidraw/excalidraw': specifier: 0.18.0-3a5ef40 version: 0.18.0-3a5ef40(@types/react-dom@18.3.1)(@types/react@18.3.12)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) @@ -561,6 +564,9 @@ importers: '@docmost/prosemirror-markdown': specifier: workspace:* version: link:../../packages/prosemirror-markdown + '@docmost/token-estimate': + specifier: workspace:* + version: link:../../packages/token-estimate '@fastify/compress': specifier: ^9.0.0 version: 9.0.0 @@ -1148,6 +1154,15 @@ importers: specifier: 4.1.6 version: 4.1.6(@opentelemetry/api@1.9.0)(@types/node@20.19.43)(@vitest/coverage-v8@4.1.6)(happy-dom@20.8.9)(jsdom@25.0.0)(vite@8.0.5(@types/node@20.19.43)(esbuild@0.28.0)(jiti@2.4.2)(less@4.2.0)(sugarss@5.0.1(postcss@8.5.14))(terser@5.39.0)(tsx@4.21.0)(yaml@2.8.3)) + packages/token-estimate: + devDependencies: + typescript: + specifier: ^5.0.0 + version: 5.9.3 + vitest: + specifier: 4.1.6 + version: 4.1.6(@opentelemetry/api@1.9.0)(@types/node@25.5.0)(@vitest/coverage-v8@4.1.6)(happy-dom@20.8.9)(jsdom@27.4.0(@noble/hashes@2.0.1))(vite@8.0.5(@types/node@25.5.0)(esbuild@0.28.0)(jiti@2.4.2)(less@4.2.0)(sugarss@5.0.1(postcss@8.5.14))(terser@5.39.0)(tsx@4.21.0)(yaml@2.8.3)) + packages: '@aashutoshrathi/word-wrap@1.2.6':