feat(ai-chat): токен-бюджет реплея истории + реактивная ветка (#490)
Вся персистентная история реплеится провайдеру КАЖДЫЙ ход, поэтому длинный чат рано или поздно упирается в контекстное окно и получает провайдерский 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) <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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/(.*)$": "<rootDir>/database/$1",
|
||||
"^@docmost/transactional/(.*)$": "<rootDir>/integrations/transactional/$1",
|
||||
"^@docmost/ee/(.*)$": "<rootDir>/ee/$1",
|
||||
"^@docmost/token-estimate$": "<rootDir>/../../../packages/token-estimate/src/index.ts",
|
||||
"^src/(.*)$": "<rootDir>/$1",
|
||||
"^@tiptap/react$": "<rootDir>/../test/stubs/tiptap-react.js"
|
||||
}
|
||||
|
||||
@@ -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<string, unknown>;
|
||||
};
|
||||
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<string, unknown>;
|
||||
};
|
||||
expect('replayOverflow' in patch.metadata).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
|
||||
@@ -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<string, unknown> | 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<string, unknown> | 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.
|
||||
|
||||
@@ -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<AiChatMessage>,
|
||||
): 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<AiChatMessage>,
|
||||
): 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
|
||||
|
||||
@@ -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<typeof makeStep>[]): 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<void>,
|
||||
): Promise<number> {
|
||||
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<typeof makeStep>[] = [];
|
||||
|
||||
// v1: each step re-persists ALL prior outputs into the trace (the O(N²) churn).
|
||||
const v1 = await walDelta(c, async () => {
|
||||
const acc: ReturnType<typeof makeStep>[] = [];
|
||||
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<typeof makeStep>[] = [];
|
||||
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<typeof makeStep>[] = [];
|
||||
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<typeof makeStep>[] = [];
|
||||
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);
|
||||
});
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -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<string, unknown>;
|
||||
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<string, unknown>;
|
||||
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<ModelMessage>,
|
||||
): 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<ModelMessage>,
|
||||
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;
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -105,6 +105,10 @@ export interface ResolvedAiConfig extends Partial<AiProviderSettings> {
|
||||
// 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
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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"]
|
||||
}
|
||||
Generated
+15
@@ -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':
|
||||
|
||||
Reference in New Issue
Block a user