perf(ai-chat): дисциплина данных — дедуп tool-outputs + токен-бюджет реплея (#490) #510
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
|
||||
@@ -13,6 +13,7 @@ import {
|
||||
compactToolOutput,
|
||||
assistantParts,
|
||||
serializeSteps,
|
||||
type StepPartsCache,
|
||||
rowToUiMessage,
|
||||
prepareAgentStep,
|
||||
stepBudgetWarning,
|
||||
@@ -28,6 +29,9 @@ import {
|
||||
FINAL_STEP_NUDGE,
|
||||
STEP_LIMIT_NO_ANSWER_MARKER,
|
||||
OUTPUT_DEGENERATION_ERROR,
|
||||
lastAssistantContextTokens,
|
||||
lastAssistantReplayOverflow,
|
||||
seedActivatedTools,
|
||||
} from './ai-chat.service';
|
||||
import type { AiChatMessage, Workspace } from '@docmost/db/types/entity.types';
|
||||
import { buildSystemPrompt } from './ai-chat.prompt';
|
||||
@@ -114,6 +118,54 @@ describe('compactToolOutput', () => {
|
||||
describe('assistantParts', () => {
|
||||
type AnyPart = Record<string, unknown>;
|
||||
|
||||
// #490 memoization: assistantParts builds each step's parts once and caches
|
||||
// them by the step OBJECT's identity, so a mid-stream flush does not
|
||||
// re-stringify every prior step's (large) output. Observable property: with a
|
||||
// shared cache, the second call over the SAME step object returns the cached
|
||||
// (identical) part array even if the step's underlying output was swapped —
|
||||
// proving the work was memoized, not redone.
|
||||
it('memoizes a step by identity (shared cache => one build per step)', () => {
|
||||
const cache: StepPartsCache = new WeakMap();
|
||||
const step = {
|
||||
text: 'x',
|
||||
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }],
|
||||
toolResults: [{ toolCallId: 'c1', toolName: 'getPage', output: { v: 1 } }],
|
||||
};
|
||||
const first = assistantParts([step], '', cache) as AnyPart[];
|
||||
expect((first.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe(
|
||||
1,
|
||||
);
|
||||
// Swap the output for a NEW value; a re-build would pick it up, a cache hit
|
||||
// keeps the first result.
|
||||
step.toolResults[0] = {
|
||||
toolCallId: 'c1',
|
||||
toolName: 'getPage',
|
||||
output: { v: 2 },
|
||||
};
|
||||
const second = assistantParts([step], '', cache) as AnyPart[];
|
||||
expect((second.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe(
|
||||
1,
|
||||
);
|
||||
// Same cached part objects are reused.
|
||||
expect(second.find((p) => p.type === 'tool-getPage')).toBe(
|
||||
first.find((p) => p.type === 'tool-getPage'),
|
||||
);
|
||||
});
|
||||
|
||||
it('without a cache, each call rebuilds (no stale memo)', () => {
|
||||
const step = {
|
||||
text: 'x',
|
||||
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }],
|
||||
toolResults: [{ toolCallId: 'c1', toolName: 'getPage', output: { v: 1 } }],
|
||||
};
|
||||
const first = assistantParts([step], '') as AnyPart[];
|
||||
step.toolResults[0].output = { v: 2 };
|
||||
const second = assistantParts([step], '') as AnyPart[];
|
||||
expect((second.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe(
|
||||
2,
|
||||
);
|
||||
});
|
||||
|
||||
it('emits output-available for a tool-call WITH a paired result', () => {
|
||||
const steps = [
|
||||
{
|
||||
@@ -231,61 +283,299 @@ describe('assistantParts', () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe('serializeSteps', () => {
|
||||
// #490 trace format v2: per call the trace stores { input } for the call and an
|
||||
// OUTCOME element — { ok: true } on success, { error, kind: 'thrown' } on a
|
||||
// thrown tool-error, { error, kind: 'interrupted' } on a mid-step abort. The tool
|
||||
// OUTPUT is no longer duplicated here (it lives once in metadata.parts).
|
||||
describe('serializeSteps (trace v2)', () => {
|
||||
it('returns null when there are no calls or results', () => {
|
||||
expect(serializeSteps([])).toBeNull();
|
||||
});
|
||||
|
||||
it('flattens calls and results into a compact trace', () => {
|
||||
it('pairs a successful call with an { ok: true } outcome and NO output', () => {
|
||||
const trace = serializeSteps([
|
||||
{
|
||||
toolCalls: [{ toolName: 'getPage', input: { id: 'p1' } }],
|
||||
toolResults: [{ toolName: 'getPage', output: { title: 'T' } }],
|
||||
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: { id: 'p1' } }],
|
||||
toolResults: [{ toolCallId: 'c1', toolName: 'getPage' }],
|
||||
},
|
||||
]) as Array<Record<string, unknown>>;
|
||||
expect(trace).toHaveLength(2);
|
||||
expect(trace[0]).toEqual({ toolName: 'getPage', input: { id: 'p1' } });
|
||||
expect(trace[1]).toEqual({ toolName: 'getPage', output: { title: 'T' } });
|
||||
expect(trace[1]).toEqual({ toolName: 'getPage', ok: true });
|
||||
// The output is NOT stored in the trace any more (dedup: it lives in parts).
|
||||
expect(trace.some((e) => 'output' in e)).toBe(false);
|
||||
});
|
||||
|
||||
it('records a THROWN tool failure (tool-error part) with its error message', () => {
|
||||
it('records a THROWN failure with { error, kind: "thrown" }', () => {
|
||||
const trace = serializeSteps([
|
||||
{
|
||||
toolCalls: [{ toolName: 'editPageText', input: { id: 'p1' } }],
|
||||
toolCalls: [
|
||||
{ toolCallId: 'c1', toolName: 'editPageText', input: { id: 'p1' } },
|
||||
],
|
||||
toolResults: [],
|
||||
content: [
|
||||
{
|
||||
type: 'tool-error',
|
||||
toolCallId: 'c1',
|
||||
toolName: 'editPageText',
|
||||
error: new Error('page is locked'),
|
||||
},
|
||||
],
|
||||
},
|
||||
]) as Array<Record<string, unknown>>;
|
||||
// The call element is followed by a paired error element (mirroring how a
|
||||
// successful result is appended), so the failure survives in the trace.
|
||||
expect(trace).toHaveLength(2);
|
||||
expect(trace[0]).toEqual({ toolName: 'editPageText', input: { id: 'p1' } });
|
||||
expect(trace[1]).toEqual({
|
||||
toolName: 'editPageText',
|
||||
error: 'page is locked',
|
||||
kind: 'thrown',
|
||||
});
|
||||
});
|
||||
|
||||
it('truncates a very long tool-error message to the tool-output limit', () => {
|
||||
it('marks an interrupted call (no result, no throw) with kind "interrupted"', () => {
|
||||
const trace = serializeSteps([
|
||||
{
|
||||
toolCalls: [
|
||||
{ toolCallId: 'c1', toolName: 'createComment', input: { x: 1 } },
|
||||
],
|
||||
toolResults: [],
|
||||
content: [],
|
||||
},
|
||||
]) as Array<Record<string, unknown>>;
|
||||
expect(trace).toHaveLength(2);
|
||||
expect(trace[1]).toEqual({
|
||||
toolName: 'createComment',
|
||||
error: 'Tool call did not complete.',
|
||||
kind: 'interrupted',
|
||||
});
|
||||
// Structurally distinct from a thrown hard-fail so it never inflates an
|
||||
// error-rate scan.
|
||||
expect((trace[1] as { kind: string }).kind).not.toBe('thrown');
|
||||
});
|
||||
|
||||
it('truncates a very long thrown-error message to the tool-output limit', () => {
|
||||
const long = 'x'.repeat(5000);
|
||||
const trace = serializeSteps([
|
||||
{
|
||||
toolCalls: [{ toolName: 'editPageText', input: {} }],
|
||||
toolCalls: [{ toolCallId: 'c1', toolName: 'editPageText', input: {} }],
|
||||
toolResults: [],
|
||||
content: [{ type: 'tool-error', toolName: 'editPageText', error: long }],
|
||||
content: [
|
||||
{
|
||||
type: 'tool-error',
|
||||
toolCallId: 'c1',
|
||||
toolName: 'editPageText',
|
||||
error: long,
|
||||
},
|
||||
],
|
||||
},
|
||||
]) as Array<Record<string, unknown>>;
|
||||
const errorText = trace[1].error as string;
|
||||
// Truncated (not the full 5000 chars) and carries the omission marker.
|
||||
expect(errorText.length).toBeLessThan(long.length);
|
||||
expect(errorText).toContain('chars omitted');
|
||||
});
|
||||
|
||||
it('pairs parallel calls in one step with their outcomes by id', () => {
|
||||
const trace = serializeSteps([
|
||||
{
|
||||
toolCalls: [
|
||||
{ toolCallId: 'a', toolName: 'getPage', input: {} },
|
||||
{ toolCallId: 'b', toolName: 'searchPages', input: {} },
|
||||
],
|
||||
toolResults: [{ toolCallId: 'b', toolName: 'searchPages' }],
|
||||
content: [
|
||||
{ type: 'tool-error', toolCallId: 'a', toolName: 'getPage', error: 'nope' },
|
||||
],
|
||||
},
|
||||
]) as Array<Record<string, unknown>>;
|
||||
// call a, outcome a (thrown), call b, outcome b (ok)
|
||||
expect(trace).toHaveLength(4);
|
||||
expect(trace[1]).toEqual({ toolName: 'getPage', error: 'nope', kind: 'thrown' });
|
||||
expect(trace[3]).toEqual({ toolName: 'searchPages', ok: true });
|
||||
});
|
||||
});
|
||||
|
||||
// #490: every assistant row flushAssistant writes carries the v2 era marker so a
|
||||
// dual-shape diagnostic query can branch on the trace shape without inspecting it.
|
||||
describe('toolTraceVersion era marker (#490)', () => {
|
||||
it('stamps metadata.toolTraceVersion = 2 on every flushed row', () => {
|
||||
const seed = flushAssistant([], '', 'streaming');
|
||||
expect(seed.metadata.toolTraceVersion).toBe(2);
|
||||
const done = flushAssistant(
|
||||
[
|
||||
{
|
||||
text: 'ok',
|
||||
toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }],
|
||||
toolResults: [{ toolCallId: 'c1', toolName: 'getPage' }],
|
||||
},
|
||||
],
|
||||
'',
|
||||
'completed',
|
||||
{ finishReason: 'stop' },
|
||||
);
|
||||
expect(done.metadata.toolTraceVersion).toBe(2);
|
||||
});
|
||||
});
|
||||
|
||||
// #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();
|
||||
});
|
||||
});
|
||||
|
||||
// #490 snapshotOpenPage fast-path: skip the full Markdown export + upsert when a
|
||||
// snapshot already exists at the page's CURRENT version (same updated_at instant).
|
||||
describe('snapshotOpenPage fast-path (#490)', () => {
|
||||
function makeSvc(existingSnapshot: unknown, pageUpdatedAt: Date) {
|
||||
const exportPageMarkdown = jest.fn(async () => '# md');
|
||||
const upsert = jest.fn(async () => undefined);
|
||||
const findByChatPage = jest.fn(async () => existingSnapshot);
|
||||
const pageRepo = {
|
||||
findById: jest.fn(async () => ({
|
||||
id: 'p1',
|
||||
workspaceId: 'ws1',
|
||||
updatedAt: pageUpdatedAt,
|
||||
})),
|
||||
};
|
||||
const svc = new AiChatService(
|
||||
{} as never, // ai
|
||||
{} as never, // aiChatRepo
|
||||
{} as never, // aiChatMessageRepo
|
||||
{ findByChatPage, upsert } as never, // aiChatPageSnapshotRepo
|
||||
{} as never, // aiSettings
|
||||
{ exportPageMarkdown } as never, // tools
|
||||
{} as never, // mcpClients
|
||||
{} as never, // aiAgentRoleRepo
|
||||
pageRepo as never, // pageRepo
|
||||
{} as never, // pageAccess
|
||||
{} as never, // environment
|
||||
);
|
||||
return { svc, exportPageMarkdown, upsert, findByChatPage };
|
||||
}
|
||||
|
||||
const args = () =>
|
||||
[
|
||||
'chat1',
|
||||
'p1',
|
||||
{ id: 'ws1' } as never,
|
||||
{ id: 'u1' } as never,
|
||||
'sess',
|
||||
] as const;
|
||||
|
||||
it('skips export + upsert when the snapshot is already at this page version', async () => {
|
||||
const t = new Date('2026-07-07T10:00:00Z');
|
||||
const { svc, exportPageMarkdown, upsert } = makeSvc(
|
||||
{ pageUpdatedAt: t, contentMd: '# md' },
|
||||
t,
|
||||
);
|
||||
await (svc as unknown as { snapshotOpenPage: (...a: unknown[]) => Promise<void> })
|
||||
.snapshotOpenPage(...args());
|
||||
expect(exportPageMarkdown).not.toHaveBeenCalled();
|
||||
expect(upsert).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('exports + upserts when the page advanced since the snapshot', async () => {
|
||||
const { svc, exportPageMarkdown, upsert } = makeSvc(
|
||||
{ pageUpdatedAt: new Date('2026-07-07T10:00:00Z'), contentMd: 'old' },
|
||||
new Date('2026-07-07T11:00:00Z'),
|
||||
);
|
||||
await (svc as unknown as { snapshotOpenPage: (...a: unknown[]) => Promise<void> })
|
||||
.snapshotOpenPage(...args());
|
||||
expect(exportPageMarkdown).toHaveBeenCalledTimes(1);
|
||||
expect(upsert).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('seeds (exports + upserts) on the first turn (no snapshot yet)', async () => {
|
||||
const { svc, exportPageMarkdown, upsert } = makeSvc(
|
||||
undefined,
|
||||
new Date('2026-07-07T10:00:00Z'),
|
||||
);
|
||||
await (svc as unknown as { snapshotOpenPage: (...a: unknown[]) => Promise<void> })
|
||||
.snapshotOpenPage(...args());
|
||||
expect(exportPageMarkdown).toHaveBeenCalledTimes(1);
|
||||
expect(upsert).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
// #490 deferred-tool activation persisted across turns.
|
||||
describe('seedActivatedTools', () => {
|
||||
const valid = new Set(['Search_web', 'getPageJson', 'diffPageVersions']);
|
||||
|
||||
it('seeds from persisted metadata, intersected with current valid names', () => {
|
||||
expect(
|
||||
seedActivatedTools(
|
||||
{ activatedTools: ['Search_web', 'getPageJson'] },
|
||||
valid,
|
||||
),
|
||||
).toEqual(['Search_web', 'getPageJson']);
|
||||
});
|
||||
|
||||
it('drops a stored tool that is no longer valid (allowlist/role changed)', () => {
|
||||
// 'Habr_publish' was activated before but is not in the current allowlist.
|
||||
expect(
|
||||
seedActivatedTools({ activatedTools: ['Search_web', 'Habr_publish'] }, valid),
|
||||
).toEqual(['Search_web']);
|
||||
});
|
||||
|
||||
it('is empty/robust for missing, non-array, or unknown-shaped metadata', () => {
|
||||
expect(seedActivatedTools(undefined, valid)).toEqual([]);
|
||||
expect(seedActivatedTools({}, valid)).toEqual([]);
|
||||
expect(seedActivatedTools({ activatedTools: 'nope' }, valid)).toEqual([]);
|
||||
expect(
|
||||
seedActivatedTools({ activatedTools: [1, 'getPageJson', null] }, valid),
|
||||
).toEqual(['getPageJson']);
|
||||
});
|
||||
|
||||
it('de-duplicates stored names', () => {
|
||||
expect(
|
||||
seedActivatedTools(
|
||||
{ activatedTools: ['getPageJson', 'getPageJson'] },
|
||||
valid,
|
||||
),
|
||||
).toEqual(['getPageJson']);
|
||||
});
|
||||
});
|
||||
|
||||
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', () => {
|
||||
@@ -618,6 +908,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
|
||||
@@ -882,6 +897,21 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
const freshPage = await this.pageRepo.findById(pageId);
|
||||
// Page deleted during the turn (or somehow foreign) => don't write.
|
||||
if (!freshPage || freshPage.workspaceId !== workspace.id) return;
|
||||
// Fast-path (#490): if a snapshot already exists at THIS page version
|
||||
// (same updated_at instant), its content is already current — skip the full
|
||||
// Markdown export + upsert entirely. A turn that did NOT touch the open page
|
||||
// (the common case) thus does no snapshot work. This mirrors the read-side
|
||||
// fast path in detectPageChange (sameInstant): both trust that a page edit
|
||||
// bumps updated_at. When the agent (or a human) DID edit the page this turn,
|
||||
// updated_at advanced, so this does not match and we re-export as before.
|
||||
const existing = await this.aiChatPageSnapshotRepo.findByChatPage(
|
||||
chatId,
|
||||
pageId,
|
||||
workspace.id,
|
||||
);
|
||||
if (existing && sameInstant(existing.pageUpdatedAt, freshPage.updatedAt)) {
|
||||
return;
|
||||
}
|
||||
const currentMd = await this.tools.exportPageMarkdown(
|
||||
user,
|
||||
sessionId,
|
||||
@@ -921,10 +951,17 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
// supplied or the supplied one does not belong to this workspace.
|
||||
let isNewChat = false;
|
||||
let chatId = body.chatId;
|
||||
// Persisted chat-level metadata bag (#490): read once here so the deferred-tool
|
||||
// activation set can be seeded from the previous turn. Undefined for a new chat.
|
||||
let chatMetadata: Record<string, unknown> | undefined;
|
||||
if (chatId) {
|
||||
const existing = await this.aiChatRepo.findById(chatId, workspace.id);
|
||||
if (!existing) {
|
||||
chatId = undefined;
|
||||
} else {
|
||||
chatMetadata = (existing.metadata ?? undefined) as
|
||||
| Record<string, unknown>
|
||||
| undefined;
|
||||
}
|
||||
}
|
||||
// The open page the client sent is attacker-controllable — BOTH its id and
|
||||
@@ -1086,7 +1123,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 +1172,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
|
||||
@@ -1326,10 +1415,19 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
// tools + ALL external MCP tools), computed from the ACTUAL toolset so an
|
||||
// external tool is loadable by its namespaced name. loadTools rejects any
|
||||
// name outside this set.
|
||||
const activatedTools = new Set<string>();
|
||||
const validDeferredNames = new Set<string>(
|
||||
Object.keys(baseTools).filter((k) => !CORE_TOOL_SET.has(k)),
|
||||
);
|
||||
// #490: seed the activation set from the chat's PERSISTED set so the model
|
||||
// does not re-run loadTools every turn to re-activate the same tools. Only
|
||||
// when deferred loading is enabled, and ALWAYS intersected with the CURRENT
|
||||
// valid deferred names — an allowlist/role change must never resurrect a tool
|
||||
// that no longer exists (prepareAgentStep would get a phantom active name).
|
||||
const activatedTools = new Set<string>(
|
||||
deferredEnabled
|
||||
? seedActivatedTools(chatMetadata, validDeferredNames)
|
||||
: [],
|
||||
);
|
||||
// Add the loadTools meta-tool ONLY when the feature is enabled; when off the
|
||||
// toolset and behavior are exactly as before.
|
||||
const tools = deferredEnabled
|
||||
@@ -1339,6 +1437,39 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
}
|
||||
: baseTools;
|
||||
|
||||
// #490: persist the (deterministically ordered) activation set back onto the
|
||||
// chat metadata at turn end, so the NEXT turn seeds from it. Once-guarded and
|
||||
// skipped when nothing new was activated (the set equals its seed) so an
|
||||
// ordinary turn adds no extra write. Preserves other metadata keys.
|
||||
let activatedToolsPersisted = false;
|
||||
const persistActivatedTools = async (): Promise<void> => {
|
||||
if (!deferredEnabled || activatedToolsPersisted || !chatId) return;
|
||||
activatedToolsPersisted = true;
|
||||
const current = [...activatedTools].sort();
|
||||
const seeded = seedActivatedTools(chatMetadata, validDeferredNames).sort();
|
||||
if (current.length === 0 || current.join(' ') === seeded.join(' ')) {
|
||||
return; // nothing new activated -> no write
|
||||
}
|
||||
try {
|
||||
await this.aiChatRepo.update(
|
||||
chatId,
|
||||
{
|
||||
metadata: {
|
||||
...(chatMetadata ?? {}),
|
||||
activatedTools: current,
|
||||
},
|
||||
} as never,
|
||||
workspace.id,
|
||||
);
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`Failed to persist activated tools (chat ${chatId}): ${
|
||||
err instanceof Error ? err.message : 'unknown error'
|
||||
}`,
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
// Accumulate the turn's streamed output so a provider error / disconnect can
|
||||
// persist the PARTIAL answer the user already saw — the SDK's onError/onAbort
|
||||
// callbacks don't hand us the in-progress text. `capturedSteps` holds finished
|
||||
@@ -1347,6 +1478,11 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
const capturedSteps: StepLike[] = [];
|
||||
let inProgressText = '';
|
||||
|
||||
// Per-turn step->parts memo (#490): shared across every flushAssistant call
|
||||
// this turn so each finished step's (large) output is JSON-stringified ONCE,
|
||||
// not re-stringified on every subsequent onStepFinish flush (was O(N²)).
|
||||
const partsCache: StepPartsCache = new WeakMap();
|
||||
|
||||
// Token-degeneration guard (#444). When the final-step lockdown is OFF, a
|
||||
// runaway repetition loop (the 255KB "loadTools." incident) is aborted via
|
||||
// this internal controller, unioned with the run/socket signal below. The
|
||||
@@ -1416,7 +1552,10 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
await this.aiChatMessageRepo.update(
|
||||
assistantId,
|
||||
workspace.id,
|
||||
flushAssistant(capturedSteps, '', 'streaming', { pageChanged }),
|
||||
flushAssistant(capturedSteps, '', 'streaming', {
|
||||
pageChanged,
|
||||
partsCache,
|
||||
}),
|
||||
{ onlyIfStreaming: true },
|
||||
);
|
||||
} catch (err) {
|
||||
@@ -1649,6 +1788,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
// closure scope here). Omitted/0 = no limit.
|
||||
maxContextTokens: resolved?.chatContextWindow,
|
||||
pageChanged,
|
||||
partsCache,
|
||||
replayTrimmedToTokens,
|
||||
}),
|
||||
);
|
||||
// #184/#487: the RUN is finalized ALWAYS (never gated on the message).
|
||||
@@ -1673,6 +1814,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
// own edits are baked in — and this also SEEDS the snapshot on the first
|
||||
// turn. Runs once across every terminal path (see snapshotTurnEnd).
|
||||
await snapshotTurnEnd();
|
||||
// #490: persist the deferred-tool activation set for the next turn.
|
||||
await persistActivatedTools();
|
||||
|
||||
// Generate the chat title for a freshly created chat AFTER the stream's
|
||||
// provider call has completed — NOT concurrently with it. The z.ai coding
|
||||
@@ -1696,7 +1839,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.
|
||||
@@ -1714,6 +1866,9 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
flushAssistant(capturedSteps, inProgressText, 'error', {
|
||||
error: errorText,
|
||||
pageChanged,
|
||||
partsCache,
|
||||
replayTrimmedToTokens,
|
||||
replayOverflow: overflow || undefined,
|
||||
}),
|
||||
);
|
||||
// #184: settle the RUN as failed, carrying the provider/transport cause.
|
||||
@@ -1723,6 +1878,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
// committed before the error must be baked into the snapshot, or the
|
||||
// next turn would mis-report it as a user edit.
|
||||
await snapshotTurnEnd();
|
||||
// #490: persist the deferred-tool activation set for the next turn.
|
||||
await persistActivatedTools();
|
||||
},
|
||||
onAbort: async ({ steps }) => {
|
||||
// #444: distinguish a degeneration abort (our internal controller) from
|
||||
@@ -1737,6 +1894,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
flushAssistant(capturedSteps, truncated, 'error', {
|
||||
error: OUTPUT_DEGENERATION_ERROR,
|
||||
pageChanged,
|
||||
partsCache,
|
||||
}),
|
||||
);
|
||||
if (runId)
|
||||
@@ -1747,6 +1905,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
);
|
||||
await closeExternalClients();
|
||||
await snapshotTurnEnd();
|
||||
// #490: persist the deferred-tool activation set for the next turn.
|
||||
await persistActivatedTools();
|
||||
return;
|
||||
}
|
||||
const partialChars =
|
||||
@@ -1771,6 +1931,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
await finalizeAssistant(
|
||||
flushAssistant(capturedSteps, inProgressText, 'aborted', {
|
||||
pageChanged,
|
||||
partsCache,
|
||||
}),
|
||||
);
|
||||
// #184: settle the RUN as aborted (an explicit user stop reached the
|
||||
@@ -1781,6 +1942,8 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
// committed before the client disconnect / stop() must be baked into the
|
||||
// snapshot, or the next turn would mis-report it as a user edit.
|
||||
await snapshotTurnEnd();
|
||||
// #490: persist the deferred-tool activation set for the next turn.
|
||||
await persistActivatedTools();
|
||||
},
|
||||
});
|
||||
|
||||
@@ -2091,6 +2254,70 @@ 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;
|
||||
}
|
||||
|
||||
/**
|
||||
* Seed the per-turn deferred-tool activation set from a chat's persisted metadata
|
||||
* (#490), INTERSECTED with the current valid deferred names. Persisting the set
|
||||
* across turns saves the model re-running loadTools every turn to re-activate the
|
||||
* same tools; intersecting on load means a changed allowlist / role can never
|
||||
* resurrect a tool that no longer exists (which would hand prepareAgentStep a
|
||||
* phantom active name). Tolerant of any stored shape — a non-array is ignored.
|
||||
*/
|
||||
export function seedActivatedTools(
|
||||
metadata: Record<string, unknown> | undefined,
|
||||
validDeferredNames: ReadonlySet<string>,
|
||||
): string[] {
|
||||
const stored = metadata?.activatedTools;
|
||||
if (!Array.isArray(stored)) return [];
|
||||
const seen = new Set<string>();
|
||||
const out: string[] = [];
|
||||
for (const name of stored) {
|
||||
if (typeof name === 'string' && validDeferredNames.has(name) && !seen.has(name)) {
|
||||
seen.add(name);
|
||||
out.push(name);
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/**
|
||||
* 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,
|
||||
@@ -2152,6 +2379,15 @@ export function sanitizeUserParts(
|
||||
/** Marker for a history row whose tool parts could not be replayed (#489). */
|
||||
export const TOOL_CONTEXT_OMITTED_MARKER = '[tool context omitted]';
|
||||
|
||||
/**
|
||||
* Synthetic error text for a tool call that neither returned a result nor threw
|
||||
* a `tool-error` — i.e. it was interrupted mid-step (an abort / server restart).
|
||||
* Shared by `assistantParts` (the replayed `output-error` part) and
|
||||
* `serializeSteps` (the `{ kind: 'interrupted' }` trace element) so the replay
|
||||
* text and the trace stay in lockstep (#490).
|
||||
*/
|
||||
export const TOOL_CALL_INCOMPLETE_TEXT = 'Tool call did not complete.';
|
||||
|
||||
/**
|
||||
* Convert persisted UI history to model messages, tolerating a single poisoned
|
||||
* row (#489). `convertToModelMessages` over the WHOLE array throws if ANY row is
|
||||
@@ -2359,71 +2595,97 @@ function normalizeToolError(error: unknown): string {
|
||||
*/
|
||||
// Exported only so the unit tests can import these pure helpers; exporting
|
||||
// them does not change runtime behavior.
|
||||
/**
|
||||
* Per-turn memo for {@link assistantParts}: a step's rebuilt parts keyed by the
|
||||
* step OBJECT's identity (#490). A finished step in `capturedSteps` keeps a stable
|
||||
* reference across every mid-stream flush, and `compactToolOutput` inside it does a
|
||||
* `JSON.stringify` of the whole (often 50–200 KB) output — so without a memo each
|
||||
* `onStepFinish` re-stringifies EVERY prior step's output (O(N²) stringify over a
|
||||
* turn). Keyed by step identity => one stringify per step per turn. WeakMap so a
|
||||
* turn's steps are GC'd with the turn.
|
||||
*/
|
||||
export type StepPartsCache = WeakMap<object, Array<Record<string, unknown>>>;
|
||||
|
||||
/** Build the parts for ONE step (text + a part per tool call). Pure. */
|
||||
function buildStepParts(step: StepLike): Array<Record<string, unknown>> {
|
||||
const parts: Array<Record<string, unknown>> = [];
|
||||
if (step.text) {
|
||||
parts.push({ type: 'text', text: step.text });
|
||||
}
|
||||
// Index this step's results by tool call id to pair calls with outputs.
|
||||
const resultsById = new Map<string, unknown>();
|
||||
for (const r of step.toolResults ?? []) {
|
||||
if (r.toolCallId) resultsById.set(r.toolCallId, r.output);
|
||||
}
|
||||
// Index this step's THROWN tool failures (ai@6 `tool-error` content parts)
|
||||
// by tool call id, so a call that failed replays with its real error text.
|
||||
const errorsById = new Map<string, unknown>();
|
||||
for (const part of step.content ?? []) {
|
||||
if (part.type === 'tool-error' && part.toolCallId) {
|
||||
errorsById.set(part.toolCallId, part.error);
|
||||
}
|
||||
}
|
||||
for (const call of step.toolCalls ?? []) {
|
||||
if (!call.toolName || !call.toolCallId) continue;
|
||||
const hasResult = resultsById.has(call.toolCallId);
|
||||
if (hasResult) {
|
||||
// output-available: the tool returned; the next turn replays its result.
|
||||
parts.push({
|
||||
type: `tool-${call.toolName}`,
|
||||
toolCallId: call.toolCallId,
|
||||
state: 'output-available',
|
||||
input: call.input,
|
||||
output: compactToolOutput(resultsById.get(call.toolCallId)),
|
||||
});
|
||||
} else if (errorsById.has(call.toolCallId)) {
|
||||
// The tool THREW: replay the REAL error so the model on the next turn
|
||||
// knows WHY the call failed (and does not blindly repeat it). An
|
||||
// output-error round-trips through convertToModelMessages as a balanced
|
||||
// tool-call + tool-result, keeping the rebuilt history valid.
|
||||
parts.push({
|
||||
type: `tool-${call.toolName}`,
|
||||
toolCallId: call.toolCallId,
|
||||
state: 'output-error',
|
||||
input: call.input,
|
||||
errorText: normalizeToolError(errorsById.get(call.toolCallId)),
|
||||
});
|
||||
} else {
|
||||
// No paired result AND no tool-error (e.g. aborted mid-step). Persisting
|
||||
// a bare tool-call (input-available) would replay as an unpaired call and
|
||||
// throw MissingToolResultsError on the next turn (convertToModelMessages
|
||||
// emits no tool-result for it). Emit a SYNTHETIC paired result instead:
|
||||
// an output-error round-trips through convertToModelMessages as a
|
||||
// balanced tool-call + tool-result, keeping the rebuilt history valid.
|
||||
parts.push({
|
||||
type: `tool-${call.toolName}`,
|
||||
toolCallId: call.toolCallId,
|
||||
state: 'output-error',
|
||||
input: call.input,
|
||||
errorText: TOOL_CALL_INCOMPLETE_TEXT,
|
||||
});
|
||||
}
|
||||
}
|
||||
return parts;
|
||||
}
|
||||
|
||||
export function assistantParts(
|
||||
steps: ReadonlyArray<StepLike> | undefined,
|
||||
fallbackText: string,
|
||||
cache?: StepPartsCache,
|
||||
): UIMessage['parts'] {
|
||||
const parts: Array<Record<string, unknown>> = [];
|
||||
let sawText = false;
|
||||
for (const step of steps ?? []) {
|
||||
if (step.text) {
|
||||
parts.push({ type: 'text', text: step.text });
|
||||
sawText = true;
|
||||
}
|
||||
// Index this step's results by tool call id to pair calls with outputs.
|
||||
const resultsById = new Map<string, unknown>();
|
||||
for (const r of step.toolResults ?? []) {
|
||||
if (r.toolCallId) resultsById.set(r.toolCallId, r.output);
|
||||
}
|
||||
// Index this step's THROWN tool failures (ai@6 `tool-error` content parts)
|
||||
// by tool call id, so a call that failed replays with its real error text.
|
||||
const errorsById = new Map<string, unknown>();
|
||||
for (const part of step.content ?? []) {
|
||||
if (part.type === 'tool-error' && part.toolCallId) {
|
||||
errorsById.set(part.toolCallId, part.error);
|
||||
}
|
||||
}
|
||||
for (const call of step.toolCalls ?? []) {
|
||||
if (!call.toolName || !call.toolCallId) continue;
|
||||
const hasResult = resultsById.has(call.toolCallId);
|
||||
if (hasResult) {
|
||||
// output-available: the tool returned; the next turn replays its result.
|
||||
parts.push({
|
||||
type: `tool-${call.toolName}`,
|
||||
toolCallId: call.toolCallId,
|
||||
state: 'output-available',
|
||||
input: call.input,
|
||||
output: compactToolOutput(resultsById.get(call.toolCallId)),
|
||||
});
|
||||
} else if (errorsById.has(call.toolCallId)) {
|
||||
// The tool THREW: replay the REAL error so the model on the next turn
|
||||
// knows WHY the call failed (and does not blindly repeat it). An
|
||||
// output-error round-trips through convertToModelMessages as a balanced
|
||||
// tool-call + tool-result, keeping the rebuilt history valid.
|
||||
parts.push({
|
||||
type: `tool-${call.toolName}`,
|
||||
toolCallId: call.toolCallId,
|
||||
state: 'output-error',
|
||||
input: call.input,
|
||||
errorText: normalizeToolError(errorsById.get(call.toolCallId)),
|
||||
});
|
||||
} else {
|
||||
// No paired result AND no tool-error (e.g. aborted mid-step). Persisting
|
||||
// a bare tool-call (input-available) would replay as an unpaired call and
|
||||
// throw MissingToolResultsError on the next turn (convertToModelMessages
|
||||
// emits no tool-result for it). Emit a SYNTHETIC paired result instead:
|
||||
// an output-error round-trips through convertToModelMessages as a
|
||||
// balanced tool-call + tool-result, keeping the rebuilt history valid.
|
||||
parts.push({
|
||||
type: `tool-${call.toolName}`,
|
||||
toolCallId: call.toolCallId,
|
||||
state: 'output-error',
|
||||
input: call.input,
|
||||
errorText: 'Tool call did not complete.',
|
||||
});
|
||||
}
|
||||
// Memoize per step object (#490): a finished step is immutable and keeps its
|
||||
// reference across flushes, so its parts (and the costly output stringify) are
|
||||
// built exactly once per turn. A cache miss (or no cache) just rebuilds.
|
||||
let stepParts = cache?.get(step as object);
|
||||
if (!stepParts) {
|
||||
stepParts = buildStepParts(step);
|
||||
cache?.set(step as object, stepParts);
|
||||
}
|
||||
parts.push(...stepParts);
|
||||
}
|
||||
const sawText = parts.some((p) => p.type === 'text');
|
||||
if (!sawText && fallbackText) {
|
||||
// No per-step text (e.g. a single final block): append the final text after
|
||||
// any tool parts so the natural call -> result -> answer order is preserved.
|
||||
@@ -2586,6 +2848,16 @@ export function flushAssistant(
|
||||
maxContextTokens?: number;
|
||||
error?: string;
|
||||
pageChanged?: { title: string; diff: string } | null;
|
||||
// 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 ?? [];
|
||||
@@ -2595,13 +2867,20 @@ export function flushAssistant(
|
||||
// in-progress step's text (the partial answer cut off by an error/abort, or
|
||||
// simply not yet flushed mid-stream) as the last text part so the persisted
|
||||
// parts match what streamed to the client.
|
||||
const parts = assistantParts(finished, '') as unknown as Array<
|
||||
Record<string, unknown>
|
||||
>;
|
||||
const parts = assistantParts(
|
||||
finished,
|
||||
'',
|
||||
extra?.partsCache,
|
||||
) as unknown as Array<Record<string, unknown>>;
|
||||
if (trailing) parts.push({ type: 'text', text: trailing });
|
||||
|
||||
const metadata: Record<string, unknown> = {
|
||||
parts: parts as unknown as UIMessage['parts'],
|
||||
// Era marker for the `tool_calls` trace shape (#490): v2 stores outcome flags
|
||||
// ({ ok } / { error, kind }) and NO tool output (the output lives once in
|
||||
// `parts`). Old rows have no marker and the legacy { output } shape; a
|
||||
// dual-shape query branches on this. Old rows are deliberately NOT migrated.
|
||||
toolTraceVersion: 2,
|
||||
};
|
||||
// finishReason: prefer an explicit one; else derive a sensible value from the
|
||||
// terminal status (so onError/onAbort records keep their historical reason).
|
||||
@@ -2617,6 +2896,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
|
||||
@@ -2642,42 +2924,85 @@ export function flushAssistant(
|
||||
|
||||
/**
|
||||
* Reduce SDK step objects to a compact, JSON-serializable trace for the
|
||||
* `tool_calls` column. Stores only what the UI action-log and history need —
|
||||
* never raw provider payloads or keys.
|
||||
* `tool_calls` column — trace format **v2** (#490).
|
||||
*
|
||||
* v2 stores, per call, ONLY the metadata a queryable trace needs — never the
|
||||
* tool OUTPUT. Before #490 each output was persisted TWICE: once here (compacted)
|
||||
* and once in `metadata.parts` (via `assistantParts`), so a 50-step run with
|
||||
* 50–200 KB outputs wrote hundreds of MB per turn (each `onStepFinish` rewrote
|
||||
* the whole row). The parts copy is the one the model replays and the UI/Markdown
|
||||
* export render, so the trace copy of the output was pure duplication. v2 keeps
|
||||
* the output ONLY in parts and reduces the trace to outcome flags.
|
||||
*
|
||||
* Element shapes (paired per call, in order):
|
||||
* - `{ toolName, input }` — the call
|
||||
* - `{ toolName, ok: true }` — it returned a result (success)
|
||||
* - `{ toolName, error, kind: 'thrown' }` — it threw a `tool-error`
|
||||
* - `{ toolName, error, kind: 'interrupted' }` — no result and no throw (an
|
||||
* abort / server restart mid-step). `kind` is MANDATORY: without it a
|
||||
* synthetic "Tool call did not complete." is indistinguishable from a real
|
||||
* hard-fail and pollutes any error-rate scan. The distinction is STRUCTURAL
|
||||
* (an `errorsById` hit vs the synthetic fallback branch), NOT a per-tool
|
||||
* classifier — soft failures stay OUT of the trace (they live in
|
||||
* `metadata.parts` outputs; a per-tool mirror would persist its own bugs).
|
||||
*
|
||||
* Rows carry `metadata.toolTraceVersion: 2` (set by {@link flushAssistant}) so a
|
||||
* dual-shape query can branch on the era. Old rows are NOT migrated (rewriting
|
||||
* giant jsonb is the very WAL churn this removes); see docs/reading-ai-logs.md.
|
||||
*/
|
||||
export function serializeSteps(
|
||||
steps: ReadonlyArray<{
|
||||
toolCalls?: ReadonlyArray<{ toolName?: string; input?: unknown }>;
|
||||
toolResults?: ReadonlyArray<{ toolName?: string; output?: unknown }>;
|
||||
toolCalls?: ReadonlyArray<{
|
||||
toolCallId?: string;
|
||||
toolName?: string;
|
||||
input?: unknown;
|
||||
}>;
|
||||
toolResults?: ReadonlyArray<{ toolCallId?: string; toolName?: string }>;
|
||||
content?: ReadonlyArray<{
|
||||
type?: string;
|
||||
toolCallId?: string;
|
||||
toolName?: string;
|
||||
error?: unknown;
|
||||
}>;
|
||||
}>,
|
||||
): unknown {
|
||||
const calls: Array<{
|
||||
toolName?: string;
|
||||
input?: unknown;
|
||||
output?: unknown;
|
||||
error?: string;
|
||||
}> = [];
|
||||
const calls: Array<
|
||||
| { toolName?: string; input?: unknown }
|
||||
| { toolName?: string; ok: true }
|
||||
| { toolName?: string; error: string; kind: 'thrown' | 'interrupted' }
|
||||
> = [];
|
||||
for (const step of steps ?? []) {
|
||||
// Index this step's results + thrown errors by tool call id, so each call is
|
||||
// paired with its outcome (mirrors assistantParts' pairing exactly).
|
||||
const resultIds = new Set<string>();
|
||||
for (const r of step.toolResults ?? []) {
|
||||
if (r.toolCallId) resultIds.add(r.toolCallId);
|
||||
}
|
||||
const errorsById = new Map<string, unknown>();
|
||||
for (const part of step.content ?? []) {
|
||||
if (part.type === 'tool-error' && part.toolCallId) {
|
||||
errorsById.set(part.toolCallId, part.error);
|
||||
}
|
||||
}
|
||||
for (const call of step.toolCalls ?? []) {
|
||||
calls.push({ toolName: call.toolName, input: call.input });
|
||||
}
|
||||
for (const r of step.toolResults ?? []) {
|
||||
calls.push({ toolName: r.toolName, output: compactToolOutput(r.output) });
|
||||
}
|
||||
// ai@6 surfaces a THROWN tool failure as a `tool-error` content part, NOT as
|
||||
// a `toolResults` entry. Record it as its own paired element (mirroring how a
|
||||
// successful result is appended) so the failure and its reason survive in the
|
||||
// trace instead of leaving an orphaned call with no result.
|
||||
for (const part of step.content ?? []) {
|
||||
if (part.type === 'tool-error') {
|
||||
if (call.toolCallId && resultIds.has(call.toolCallId)) {
|
||||
// Success: the output itself lives in metadata.parts, not here.
|
||||
calls.push({ toolName: call.toolName, ok: true });
|
||||
} else if (call.toolCallId && errorsById.has(call.toolCallId)) {
|
||||
// Hard fail: the tool threw. Persist the real (bounded) reason.
|
||||
calls.push({
|
||||
toolName: part.toolName,
|
||||
error: normalizeToolError(part.error),
|
||||
toolName: call.toolName,
|
||||
error: normalizeToolError(errorsById.get(call.toolCallId)),
|
||||
kind: 'thrown',
|
||||
});
|
||||
} else {
|
||||
// Neither a result nor a throw: interrupted mid-step (abort/restart).
|
||||
// Marked structurally so it never inflates a thrown-error count.
|
||||
calls.push({
|
||||
toolName: call.toolName,
|
||||
error: TOOL_CALL_INCOMPLETE_TEXT,
|
||||
kind: 'interrupted',
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
import { type Kysely, sql } from 'kysely';
|
||||
|
||||
export async function up(db: Kysely<any>): Promise<void> {
|
||||
// Chat-level metadata bag (#490). First use: the deferred-tool ACTIVATION set
|
||||
// (`activatedTools`) is persisted here so it survives across turns — previously
|
||||
// the set was reset every turn, forcing the model to re-run loadTools and pay a
|
||||
// fresh round-trip to re-activate the same tools each turn. On load the stored
|
||||
// set is intersected with the current valid deferred names, so an allowlist /
|
||||
// role change can never inject a now-nonexistent tool.
|
||||
//
|
||||
// jsonb, defaulted to '{}' so every row (incl. pre-migration ones, backfilled
|
||||
// by the default) is a readable object — the app never has to null-guard the
|
||||
// bag itself, only individual keys.
|
||||
await db.schema
|
||||
.alterTable('ai_chats')
|
||||
.addColumn('metadata', 'jsonb', (col) =>
|
||||
col.notNull().defaultTo(sql`'{}'::jsonb`),
|
||||
)
|
||||
.execute();
|
||||
}
|
||||
|
||||
export async function down(db: Kysely<any>): Promise<void> {
|
||||
await db.schema.alterTable('ai_chats').dropColumn('metadata').execute();
|
||||
}
|
||||
+3
@@ -606,6 +606,9 @@ export interface AiChats {
|
||||
// The document the chat was created in (open page at first message). NULL =>
|
||||
// started outside any document. ON DELETE SET NULL on the page FK.
|
||||
pageId: string | null;
|
||||
// Chat-level metadata bag (#490). jsonb, defaulted to '{}'. First key:
|
||||
// `activatedTools` — the deferred-tool activation set persisted across turns.
|
||||
metadata: Generated<Json>;
|
||||
createdAt: Generated<Timestamp>;
|
||||
updatedAt: Generated<Timestamp>;
|
||||
deletedAt: Timestamp | null;
|
||||
|
||||
@@ -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
|
||||
|
||||
+174
-78
@@ -8,19 +8,35 @@ real pain (a "which tools fail most?" analysis that confidently answered
|
||||
|
||||
Read the **Gotchas** section before you trust any error count.
|
||||
|
||||
> **TWO ERAS — check the marker first.** The `tool_calls` shape changed in **#490
|
||||
> (trace v2)**. A row written by v2 carries `metadata.toolTraceVersion = 2`; older
|
||||
> rows have no such key. The two shapes store DIFFERENT things (v2 dropped the tool
|
||||
> OUTPUT from the trace), so **every query below is dual-shape** — branch on the
|
||||
> marker. **Never compare an aggregate or trend across the era boundary**: a metric
|
||||
> jump on the cut-over week is an artifact of the shape change, not a behavior
|
||||
> change.
|
||||
|
||||
## TL;DR
|
||||
|
||||
- Agent chats live in Postgres, DB `docmost`, tables `ai_chat_*`.
|
||||
- Each tool invocation is stored as **two** array elements (a `tool-call` part and
|
||||
a `tool-result` part), so naive counting double-counts.
|
||||
- **A tool that *throws* writes no result part.** Since the #407 fix its error is
|
||||
persisted as a dedicated `{toolName, error}` element in `tool_calls` (queryable +
|
||||
replayed to the model). **Rows written before #407 still drop it** — the error is
|
||||
nowhere in the DB and shows only in the live UI. So `isError` / `success=false`
|
||||
scans under-report by design, and pre-#407 thrown errors are invisible.
|
||||
- To find where agents fail: (1) soft-failure markers in `tool_calls`, (2) the new
|
||||
`error` field for thrown errors (new rows) / the orphan-gap proxy (old rows),
|
||||
(3) server logs / the live UI for full stack traces beyond the truncated message.
|
||||
- **Era marker:** `metadata.toolTraceVersion = 2` ⇒ v2 (#490) row; absent ⇒ legacy row.
|
||||
- Each tool invocation is stored as **two** consecutive array elements — a
|
||||
`tool-call` part then an OUTCOME part — so naive counting double-counts.
|
||||
- **v2 (#490):** outcome is `{toolName, ok: true}` on success, or
|
||||
`{toolName, error, kind: 'thrown'|'interrupted'}` on failure. The tool **OUTPUT
|
||||
is NOT in `tool_calls`** any more — it lives once in `metadata.parts` (this
|
||||
removed a hundreds-of-MB-per-run write duplication). Soft-failure analysis
|
||||
therefore reads `metadata.parts`, not `tool_calls`.
|
||||
- **legacy:** outcome is `{toolName, output}` on success; a **thrown** failure is
|
||||
a `{toolName, error}` element **only on rows after #407**, and is dropped
|
||||
entirely (silent orphan) on pre-#407 rows.
|
||||
- **A tool that *throws* writes no result part.** In v2 it is a
|
||||
`{error, kind:'thrown'}` element; an interrupted/aborted call is a distinct
|
||||
`{error, kind:'interrupted'}`. `isError`/`success=false` scans read the *output*
|
||||
and so under-report thrown failures in every era.
|
||||
- To find where agents fail: (1) soft-failure markers in `metadata.parts` outputs
|
||||
(v2) / `tool_calls` outputs (legacy), (2) the `error`/`kind` fields for thrown
|
||||
failures (v2 + post-#407), (3) server logs / the live UI for full stack traces.
|
||||
|
||||
## Where the data lives
|
||||
|
||||
@@ -53,33 +69,67 @@ are rows in `workspaces`, not separate deployments.
|
||||
separate `tool` role), `content` (text), `tool_calls` (jsonb array), `metadata`
|
||||
(jsonb, holds run `error` + rendered `parts`), `status`, `tsv` (full-text index).
|
||||
|
||||
## Era marker — check this before every query
|
||||
|
||||
```sql
|
||||
-- how many rows are in each era?
|
||||
SELECT COALESCE((metadata->>'toolTraceVersion'), 'legacy') AS era, count(*)
|
||||
FROM ai_chat_messages
|
||||
WHERE role = 'assistant' AND jsonb_typeof(tool_calls) = 'array'
|
||||
GROUP BY 1 ORDER BY 2 DESC;
|
||||
```
|
||||
|
||||
- `toolTraceVersion = '2'` → **v2** (#490): outcome flags, **no output in the trace**.
|
||||
- `NULL` (`'legacy'`) → pre-#490: outcome carries the tool `output` inline.
|
||||
|
||||
**Do not trend a metric across the cut-over.** The shape change alone shifts counts
|
||||
(e.g. "elements with `output`" collapses to zero for v2), so a week that straddles
|
||||
the boundary shows an artifact, not a behavior change. Segment by era, or restrict to
|
||||
one era, before comparing.
|
||||
|
||||
## How tool calls are stored — READ THIS
|
||||
|
||||
Tool calls are **not** one-object-per-call. Each logical invocation is split into
|
||||
two consecutive elements of the `tool_calls` array:
|
||||
two consecutive elements of the `tool_calls` array — a **call** then an **outcome**.
|
||||
The outcome shape is era-dependent:
|
||||
|
||||
```text
|
||||
index 0: { "toolName": "getPage", "input": { "pageId": "…" } } ← tool-call (has input, NO output)
|
||||
index 1: { "toolName": "getPage", "output": { … } } ← tool-result (has output, NO input)
|
||||
# v2 (#490) — metadata.toolTraceVersion = 2
|
||||
index 0: { "toolName":"getPage", "input":{...} } ← call (has input)
|
||||
index 1: { "toolName":"getPage", "ok":true } ← success (NO output here)
|
||||
or : { "toolName":"getPage", "error":"…", "kind":"thrown" } ← threw
|
||||
or : { "toolName":"getPage", "error":"…", "kind":"interrupted" } ← aborted mid-step
|
||||
|
||||
# legacy — no toolTraceVersion
|
||||
index 0: { "toolName":"getPage", "input":{...} } ← call (has input, NO output)
|
||||
index 1: { "toolName":"getPage", "output":{...} } ← success (has output)
|
||||
or : { "toolName":"getPage", "error":"…" } ← threw (post-#407 only)
|
||||
```
|
||||
|
||||
The keys that appear on an element are `toolName`, `input`, `output`, and — for a
|
||||
**thrown** failure on rows written after the #407 fix — `error` (the tool's error
|
||||
message; see the "Hard failures" section below). There is no `state`, no `errorText`,
|
||||
no `type`. On pre-#407 rows a thrown failure has NO paired result element at all
|
||||
(silent orphan). Consequences:
|
||||
The keys that can appear: `toolName`, `input` (call), and on the outcome — **v2:**
|
||||
`ok` **or** `error`+`kind`; **legacy:** `output` **or** (post-#407) `error`. There is
|
||||
no `state`, no `errorText`, no `type` in `tool_calls` (those live on `metadata.parts`).
|
||||
Consequences:
|
||||
|
||||
1. **Real invocation count = elements that have `output` or `error`.** Counting every
|
||||
element double-counts (you get ~2× and a spurious "~50% of every tool has no output").
|
||||
2. **Pairing:** a call = a `tool-call` part followed by its result part. A success
|
||||
carries `output`; a thrown failure (post-#407) carries `error` instead. Both carry
|
||||
`toolName`, so you can group by tool on either.
|
||||
1. **Real invocation count** — count the OUTCOME elements, not every element (else you
|
||||
double-count): **v2** = elements with `ok` or `error`; **legacy** = elements with
|
||||
`output` or `error`.
|
||||
2. **Pairing:** a call (`input`) is followed by its outcome. `toolName` is on both, so
|
||||
you can group by tool on either. In v2 the `kind` field separates a real hard-fail
|
||||
(`thrown`) from an aborted call (`interrupted`) — a distinction legacy rows cannot
|
||||
make (both are orphans; see below).
|
||||
3. **The tool OUTPUT is only in `metadata.parts` on v2 rows.** To inspect what a tool
|
||||
returned (soft-error markers, page bodies) on a v2 row, read the parts
|
||||
(`part->>'type' LIKE 'tool-%'`, `part->>'state' = 'output-available'`, `part->'output'`),
|
||||
not `tool_calls`.
|
||||
|
||||
## The two classes of failure (and which the DB can see)
|
||||
|
||||
### 1. Soft failures — tool RAN and returned an error-shaped result → PERSISTED ✅
|
||||
|
||||
These are visible in the `tool-result` `output`. The marker differs per tool:
|
||||
These are visible in the tool `output` — **on v2 rows in `metadata.parts`** (the
|
||||
`output-available` part's `output`), on **legacy rows in the `tool_calls` outcome
|
||||
element's `output`**. The marker differs per tool:
|
||||
|
||||
| Tool(s) | Error marker in `output` |
|
||||
| --- | --- |
|
||||
@@ -91,37 +141,32 @@ These are visible in the `tool-result` `output`. The marker differs per tool:
|
||||
Note `editPageText` returns `failed: []` on success — filtering on the *presence*
|
||||
of the key gives false positives; filter on **non-empty**.
|
||||
|
||||
### 2. Hard failures — tool THREW → NOW PERSISTED ✅ (since the #407 fix)
|
||||
### 2. Hard failures — tool THREW → PERSISTED ✅
|
||||
|
||||
When a tool throws (the classic one is `patchNode` / `insertNode` / `tableUpdateCell`
|
||||
→ `Failed to encode document to Yjs (fromJSON): Unknown node type: undefined`), the
|
||||
runtime still writes **no `tool-result` part** — the failure is an ai@6 `tool-error`
|
||||
content part instead. **Since the #407 fix, that error is persisted**: `serializeSteps`
|
||||
appends a dedicated element `{toolName, error: "<message>"}` right after the failed
|
||||
call, mirroring how a successful `{toolName, output}` element is appended. So a thrown
|
||||
error now leaves a queryable `error` field carrying its (truncated) reason, and the
|
||||
same real text is replayed to the model on the next turn (an `output-error` part with
|
||||
the real `errorText`, no longer the `'Tool call did not complete.'` placeholder).
|
||||
runtime writes **no `tool-result` part** — the failure is an ai@6 `tool-error` content
|
||||
part. How that lands in `tool_calls` depends on the era:
|
||||
|
||||
**Cutover caveat — old rows keep the old blind shape.** Rows written **before** this
|
||||
change have the two-part shape (`call` + `output` only) and simply **drop** thrown
|
||||
errors, leaving a silent **orphan** (a `call` with no `output` *and* no `error`). Rows
|
||||
written **after** the fix additionally carry the `error` element. So:
|
||||
- **v2 (#490):** a `{toolName, error, kind:'thrown'}` outcome element. An interrupted /
|
||||
aborted mid-step call is a **distinct** `{toolName, error:'Tool call did not
|
||||
complete.', kind:'interrupted'}` element — so you can tell a real hard-fail from an
|
||||
abort **directly, without the orphan heuristic**. Query `kind = 'thrown'`.
|
||||
- **post-#407 legacy:** a `{toolName, error}` element (no `kind`) right after the call.
|
||||
- **pre-#407 legacy:** the error is **dropped** — a silent **orphan** (a `call` with no
|
||||
`output` *and* no `error`).
|
||||
|
||||
- **New rows:** query the `error` field directly (see the hard-error query below) — no
|
||||
orphan heuristic needed for thrown failures.
|
||||
- **Old rows (pre-#407):** the only DB-side proxy is still an **orphan**: a `tool-call`
|
||||
part with no matching `tool-result` *and* no `error`. Orphans also appear when a run
|
||||
is **aborted** mid-flight (server restart), so a high-volume tool (`createComment`,
|
||||
`searchInPage`, `Search_web_search`) shows orphans from aborts, not real errors on
|
||||
old rows. Treat the orphan gap as an *upper bound*, and cross-check the tool: a gap on
|
||||
a structural editor (`patchNode`, `insertNode`, `updatePageJson`, `transformPage`) is
|
||||
almost certainly a thrown Yjs-encode error; a gap on `createComment` is mostly aborts.
|
||||
The same real error text is replayed to the model on the next turn (an `output-error`
|
||||
part with the real `errorText`, from `metadata.parts`), in every era.
|
||||
|
||||
A note on the aborted-call fallback: a call with **neither** a result **nor** a
|
||||
`tool-error` (genuinely interrupted mid-step) still replays with the
|
||||
`'Tool call did not complete.'` placeholder and persists as an orphan — that path is
|
||||
unchanged, and is distinct from a real thrown error, which now carries `error`.
|
||||
**Cutover caveat.** Only pre-#407 legacy rows need the orphan proxy: an orphan is a
|
||||
`tool-call` with no matching outcome. Orphans there also appear when a run is **aborted**
|
||||
mid-flight (server restart), so a high-volume tool (`createComment`, `searchInPage`,
|
||||
`Search_web_search`) shows orphans from aborts, not real errors. Treat the orphan gap as
|
||||
an *upper bound* and cross-check the tool: a gap on a structural editor (`patchNode`,
|
||||
`insertNode`, `updatePageJson`, `transformPage`) is almost certainly a thrown Yjs-encode
|
||||
error; a gap on `createComment` is mostly aborts. **On v2 rows this ambiguity is gone**
|
||||
— `kind` labels each outcome.
|
||||
|
||||
### 3. Run-level failures → `ai_chat_runs`
|
||||
|
||||
@@ -134,22 +179,34 @@ the wild: `Run interrupted by a server restart.` (aborts) and
|
||||
|
||||
Run all of these via `docker exec gitmost-postgresql psql -U docmost -d docmost -P pager=off -c "…"`.
|
||||
|
||||
**Real invocation count per tool** (result parts only — the correct denominator):
|
||||
**Real invocation count per tool** (outcome parts only — the correct denominator).
|
||||
Dual-shape: a v2 outcome has `ok` or `error`; a legacy outcome has `output` or `error`:
|
||||
|
||||
```sql
|
||||
SELECT elem->>'toolName' AS tool, count(*) AS calls
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array' AND elem ? 'output'
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array'
|
||||
AND (elem ? 'ok' OR elem ? 'output' OR elem ? 'error')
|
||||
GROUP BY 1 ORDER BY 2 DESC;
|
||||
```
|
||||
|
||||
**Soft errors per tool** (everything the DB can honestly see):
|
||||
**Soft errors per tool.** The soft-error marker lives in the tool OUTPUT — which on
|
||||
**v2 rows is in `metadata.parts`**, on **legacy rows is in the `tool_calls` outcome
|
||||
element**. This query UNIONs both eras, projecting each output as `o`:
|
||||
|
||||
```sql
|
||||
WITH res AS (
|
||||
-- v2 (#490): output is in metadata.parts (output-available tool parts)
|
||||
SELECT part->>'type' AS tool, part->'output' AS o
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.metadata->'parts') part
|
||||
WHERE (m.metadata->>'toolTraceVersion') = '2'
|
||||
AND part->>'type' LIKE 'tool-%' AND part->>'state' = 'output-available'
|
||||
UNION ALL
|
||||
-- legacy: output is inline in the tool_calls outcome element
|
||||
SELECT elem->>'toolName' AS tool, elem->'output' AS o
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array' AND elem ? 'output'
|
||||
WHERE (m.metadata->>'toolTraceVersion') IS NULL
|
||||
AND jsonb_typeof(m.tool_calls) = 'array' AND elem ? 'output'
|
||||
)
|
||||
SELECT tool, count(*) AS calls,
|
||||
sum(COALESCE(
|
||||
@@ -167,13 +224,23 @@ FROM res GROUP BY tool HAVING sum(COALESCE(
|
||||
ORDER BY soft_errors DESC;
|
||||
```
|
||||
|
||||
**`editPageText` failure reasons** (the most common real agent mistake — bad `find`):
|
||||
Note the v2 `tool` label is the part type (`tool-editPageText`); strip the `tool-`
|
||||
prefix if you join it against the legacy `toolName`.
|
||||
|
||||
**`editPageText` failure reasons** (the most common real agent mistake — bad `find`).
|
||||
Same dual-shape output source:
|
||||
|
||||
```sql
|
||||
WITH res AS (
|
||||
SELECT part->'output' AS o
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.metadata->'parts') part
|
||||
WHERE (m.metadata->>'toolTraceVersion') = '2'
|
||||
AND part->>'type' = 'tool-editPageText' AND part->>'state' = 'output-available'
|
||||
UNION ALL
|
||||
SELECT elem->'output' AS o
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array'
|
||||
WHERE (m.metadata->>'toolTraceVersion') IS NULL
|
||||
AND jsonb_typeof(m.tool_calls) = 'array'
|
||||
AND elem->>'toolName' = 'editPageText' AND elem ? 'output'
|
||||
)
|
||||
SELECT f->>'reason' AS reason, count(*)
|
||||
@@ -182,30 +249,43 @@ WHERE jsonb_typeof(o->'failed') = 'array'
|
||||
GROUP BY 1 ORDER BY 2 DESC;
|
||||
```
|
||||
|
||||
**Hard errors — persisted `error` field per tool (NEW rows, since #407)** — thrown
|
||||
tool failures now carry their real reason, so query them directly:
|
||||
**Hard errors — persisted `error` field per tool (v2 + post-#407 rows)** — thrown tool
|
||||
failures carry their real reason, so query them directly. On **v2** rows exclude the
|
||||
`interrupted` kind so an aborted call is not counted as a hard-fail:
|
||||
|
||||
```sql
|
||||
SELECT elem->>'toolName' AS tool, count(*) AS thrown_errors,
|
||||
min(elem->>'error') AS sample_error
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array' AND elem ? 'error'
|
||||
-- v2 rows label the kind; a legacy error element has no kind (count it).
|
||||
AND COALESCE(elem->>'kind', 'thrown') = 'thrown'
|
||||
GROUP BY 1 ORDER BY 2 DESC;
|
||||
```
|
||||
|
||||
Aborted mid-step calls on v2 rows are a distinct, directly countable population:
|
||||
|
||||
```sql
|
||||
SELECT elem->>'toolName' AS tool, count(*) AS interrupted
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array' AND elem->>'kind' = 'interrupted'
|
||||
GROUP BY 1 ORDER BY 2 DESC;
|
||||
```
|
||||
|
||||
**Hard-error proxy for OLD rows (pre-#407) — orphan gap per tool, WITH a spread column**
|
||||
(call parts minus result parts, plus how many distinct chats the gap is spread across).
|
||||
This covers rows written before thrown errors were persisted; on new rows a thrown
|
||||
failure now has its own `error` element (use the query above) and an orphan means only
|
||||
a genuinely aborted mid-step call:
|
||||
(call parts minus outcome parts, plus how many distinct chats the gap is spread across).
|
||||
This is needed ONLY for pre-#407 legacy rows (v2 and post-#407 rows carry the error /
|
||||
`kind` directly — use the queries above). The `WHERE` restricts to the legacy era so v2
|
||||
rows (where an `ok` outcome is not an `output`) never produce phantom orphans:
|
||||
|
||||
```sql
|
||||
WITH parts AS (
|
||||
SELECT m.chat_id, elem->>'toolName' AS tool,
|
||||
(elem ? 'input' AND NOT (elem ? 'output')) AS is_call,
|
||||
(elem ? 'output' OR elem ? 'error') AS is_result
|
||||
(elem ? 'input' AND NOT (elem ? 'output') AND NOT (elem ? 'ok')) AS is_call,
|
||||
(elem ? 'output' OR elem ? 'error' OR elem ? 'ok') AS is_result
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
WHERE jsonb_typeof(m.tool_calls) = 'array' AND m.role = 'assistant'
|
||||
AND (m.metadata->>'toolTraceVersion') IS NULL
|
||||
),
|
||||
per_chat AS (
|
||||
SELECT tool, chat_id, sum(is_call::int) - sum(is_result::int) AS gap
|
||||
@@ -261,11 +341,21 @@ WHERE tsv @@ websearch_to_tsquery('english', 'some phrase') LIMIT 20;
|
||||
|
||||
## Don't blow up your context
|
||||
|
||||
A single `tool_calls` row can be **300–400 KB** (results embed full page content and
|
||||
search payloads). Never `SELECT tool_calls` (or `jsonb_pretty(tool_calls)`) raw.
|
||||
Always project just the keys you need and truncate:
|
||||
Tool outputs embed full page content and search payloads (hundreds of KB per row).
|
||||
On **legacy** rows they are in `tool_calls`; on **v2** rows they moved to
|
||||
`metadata->'parts'` (the `tool_calls` trace itself is now small). Never `SELECT
|
||||
tool_calls` / `metadata` (or `jsonb_pretty(...)`) raw — project just the keys you need
|
||||
and truncate:
|
||||
|
||||
```sql
|
||||
-- v2: outputs live in metadata.parts
|
||||
SELECT part->>'type',
|
||||
left(regexp_replace((part->'output')::text, '\s+', ' ', 'g'), 200)
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.metadata->'parts') part
|
||||
WHERE (m.metadata->>'toolTraceVersion') = '2'
|
||||
AND part->>'state' = 'output-available' LIMIT 5;
|
||||
|
||||
-- legacy: outputs live in tool_calls
|
||||
SELECT elem->>'toolName',
|
||||
left(regexp_replace((elem->'output')::text, '\s+', ' ', 'g'), 200)
|
||||
FROM ai_chat_messages m, jsonb_array_elements(m.tool_calls) elem
|
||||
@@ -280,26 +370,32 @@ docker compose -p gitmost logs -f --tail=100 # whole stack
|
||||
```
|
||||
|
||||
Logging is `json-file`, `max-size=10m max-file=5` → ~50 MB retained, then rotated,
|
||||
and **wiped on container recreate**. Since the #407 fix, thrown-tool error text is
|
||||
**persisted in the `error` field** of `tool_calls` (see the hard-error query above), so
|
||||
you no longer depend on live logs for it. Logs/live UI remain useful for **pre-#407
|
||||
rows** (whose thrown errors were dropped) and for full stack traces beyond the
|
||||
truncated stored message. A per-tool `tool_calls_total{tool,status}` metric to
|
||||
VictoriaMetrics is still a possible future add for aggregate dashboards.
|
||||
and **wiped on container recreate**. Thrown-tool error text is **persisted** — in the
|
||||
`error` field of `tool_calls` (v2 `kind:'thrown'` / post-#407 legacy) — so you no longer
|
||||
depend on live logs for it. Logs/live UI remain useful for **pre-#407 rows** (whose
|
||||
thrown errors were dropped) and for full stack traces beyond the truncated stored
|
||||
message. A per-tool `tool_calls_total{tool,status}` metric to VictoriaMetrics is still a
|
||||
possible future add for aggregate dashboards.
|
||||
|
||||
## Gotchas checklist
|
||||
|
||||
- [ ] Counting every `tool_calls` element → **overcount**. Count `output` elements; add `error` elements for thrown failures (new rows), but don't count both as invocations.
|
||||
- [ ] `isError` / `success=false` ≈ 0 does **not** mean "no errors" — thrown errors are a separate `error` element (new rows) or dropped entirely (pre-#407 rows).
|
||||
- [ ] Thrown errors persist only on rows written **after the #407 fix** — pre-#407 rows still drop them (orphan only). Mind the cutover when trending over time.
|
||||
- [ ] **Check `metadata.toolTraceVersion` first.** v2 (`= 2`) has no output in `tool_calls`; legacy has it inline. Never trend a metric across the era boundary.
|
||||
- [ ] Counting every `tool_calls` element → **overcount**. Count OUTCOME elements — v2: `ok` or `error`; legacy: `output` or `error` — never both call+outcome as invocations.
|
||||
- [ ] `isError` / `success=false` ≈ 0 does **not** mean "no errors" — thrown errors are an `error` element (v2 `kind:'thrown'` / post-#407), not in the output.
|
||||
- [ ] **v2:** soft-error markers (the tool output) are in `metadata.parts`, NOT `tool_calls`. Legacy: they are in the `tool_calls` outcome `output`.
|
||||
- [ ] **v2:** `kind` splits a real hard-fail (`thrown`) from an aborted call (`interrupted`) directly — no orphan heuristic needed. The orphan gap is a pre-#407-legacy-only proxy.
|
||||
- [ ] `editPageText.failed` is `[]` on success — test for **non-empty**, not presence.
|
||||
- [ ] Orphan gap on OLD rows mixes thrown errors **and** aborted runs — split by tool. On NEW rows a thrown error is its own `error` element, so a gap ≈ aborted call.
|
||||
- [ ] `aborted` runs = server restarts, `failed` runs = provider overload — not agent mistakes.
|
||||
- [ ] Never dump a raw `tool_calls` cell — it can be hundreds of KB.
|
||||
- [ ] Logs are ephemeral (≤50 MB, wiped on recreate) — grab hard-error text live.
|
||||
- [ ] Never dump a raw `tool_calls` **or** `metadata.parts` cell — outputs are hundreds of KB.
|
||||
- [ ] Logs are ephemeral (≤50 MB, wiped on recreate) — grab pre-#407 hard-error text live.
|
||||
|
||||
## Snapshot (2026-07-07, illustrative — rerun the queries for current numbers)
|
||||
|
||||
> All rows in this snapshot predate #490, so they are **legacy-era** (outputs inline in
|
||||
> `tool_calls`, orphan proxy for thrown errors). Do not trend these numbers against v2
|
||||
> rows — segment by `toolTraceVersion` first.
|
||||
|
||||
|
||||
- 226 chats, 732 messages, 46 runs; ~4 400 real tool invocations.
|
||||
- Soft errors (persisted): `editPageText` 4/79 (bad/non-unique `find`) + 9 markdown-in-`find` warnings; `semanticSearch` 3/4 (`unavailable`); `Habr_update_draft_from_docmost` 1/2 (`doc` sent as object, not string).
|
||||
- Missing-result proxy, read WITH the spread column:
|
||||
|
||||
@@ -59,6 +59,39 @@ import {
|
||||
mergeFootnoteDefinitions,
|
||||
} from "../lib/transforms.js";
|
||||
|
||||
// Max concurrent per-page comment fetches in checkNewComments (#490). The scan is
|
||||
// O(N) independent REST reads over the working set; running them one-at-a-time made
|
||||
// a large space linear in round-trips. A small cap parallelizes without hammering
|
||||
// the server (or exhausting sockets). 6 is a conservative middle of the 5–8 band.
|
||||
const COMMENT_SCAN_CONCURRENCY = 6;
|
||||
|
||||
/**
|
||||
* Map `items` through `fn` with at most `limit` in flight, preserving INPUT ORDER
|
||||
* in the returned array. A tiny bounded pool (no p-limit dependency): `limit`
|
||||
* workers pull the next index off a shared cursor until the list is drained.
|
||||
*/
|
||||
async function mapWithConcurrency<T, R>(
|
||||
items: readonly T[],
|
||||
limit: number,
|
||||
fn: (item: T, index: number) => Promise<R>,
|
||||
): Promise<R[]> {
|
||||
const results = new Array<R>(items.length);
|
||||
let cursor = 0;
|
||||
const worker = async (): Promise<void> => {
|
||||
for (;;) {
|
||||
const i = cursor++;
|
||||
if (i >= items.length) return;
|
||||
results[i] = await fn(items[i], i);
|
||||
}
|
||||
};
|
||||
const workers = Array.from(
|
||||
{ length: Math.max(1, Math.min(limit, items.length)) },
|
||||
() => worker(),
|
||||
);
|
||||
await Promise.all(workers);
|
||||
return results;
|
||||
}
|
||||
|
||||
// Public method surface of CommentsMixin (issue #450) — a NAMED type so the factory
|
||||
// return type is expressible in the emitted .d.ts (the anonymous mixin class
|
||||
// carries the base's protected shared state, which would otherwise trip TS4094).
|
||||
@@ -655,27 +688,32 @@ export function CommentsMixin<TBase extends GConstructor<DocmostClientContext>>(
|
||||
parentPageId,
|
||||
);
|
||||
|
||||
// 2. Fetch comments for each page, keep ones created after since
|
||||
const results: any[] = [];
|
||||
for (const page of pagesInScope) {
|
||||
try {
|
||||
// Full feed (incl. resolved): a "new comments since" scan reports all
|
||||
// recent activity; the active-only filter is scoped to listComments.
|
||||
const comments = (await this.listComments(page.id, true)).items;
|
||||
const newComments = comments.filter(
|
||||
(c: any) => new Date(c.createdAt) > sinceDate,
|
||||
);
|
||||
if (newComments.length > 0) {
|
||||
results.push({
|
||||
pageId: page.id,
|
||||
pageTitle: page.title,
|
||||
comments: newComments,
|
||||
});
|
||||
// 2. Fetch comments for each page, keep ones created after since. Runs with
|
||||
// bounded concurrency (#490) instead of one-at-a-time — the per-page reads are
|
||||
// independent, so a large working set no longer costs O(N) serial round-trips.
|
||||
// Order is preserved (mapWithConcurrency keeps input order), so the output is
|
||||
// deterministic regardless of which fetch finishes first.
|
||||
const perPage = await mapWithConcurrency(
|
||||
pagesInScope,
|
||||
COMMENT_SCAN_CONCURRENCY,
|
||||
async (page: any) => {
|
||||
try {
|
||||
// Full feed (incl. resolved): a "new comments since" scan reports all
|
||||
// recent activity; the active-only filter is scoped to listComments.
|
||||
const comments = (await this.listComments(page.id, true)).items;
|
||||
const newComments = comments.filter(
|
||||
(c: any) => new Date(c.createdAt) > sinceDate,
|
||||
);
|
||||
return newComments.length > 0
|
||||
? { pageId: page.id, pageTitle: page.title, comments: newComments }
|
||||
: null;
|
||||
} catch (e: any) {
|
||||
// Skip pages with errors (e.g. deleted between calls)
|
||||
return null;
|
||||
}
|
||||
} catch (e: any) {
|
||||
// Skip pages with errors (e.g. deleted between calls)
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
const results: any[] = perPage.filter((r): r is any => r !== null);
|
||||
|
||||
const totalNewComments = results.reduce(
|
||||
(sum, r) => sum + r.comments.length,
|
||||
|
||||
@@ -442,3 +442,71 @@ test("checkNewComments subtree includes the root without a separate getPageRaw",
|
||||
assert.equal(result.checkedPages, 2, "root + one descendant scanned");
|
||||
assert.equal(result.totalNewComments, 1, "the root's fresh comment found");
|
||||
});
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// 6) checkNewComments parallelism (#490): the per-page comment fetches run with
|
||||
// bounded concurrency (not one-at-a-time), and the results still preserve the
|
||||
// page order deterministically regardless of which fetch finishes first.
|
||||
// -----------------------------------------------------------------------------
|
||||
test("checkNewComments fetches pages concurrently (bounded) and preserves order", async () => {
|
||||
// A subtree with 12 descendants so the scan has plenty to parallelize.
|
||||
const NODES = [{ id: "parent", title: "Parent", parentPageId: null, hasChildren: true }];
|
||||
for (let i = 0; i < 12; i++) {
|
||||
NODES.push({ id: `k${i}`, title: `Kid ${i}`, parentPageId: "parent", hasChildren: false });
|
||||
}
|
||||
|
||||
let inFlight = 0;
|
||||
let maxInFlight = 0;
|
||||
|
||||
const { baseURL } = await spawn(async (req, res) => {
|
||||
const raw = await readBody(req);
|
||||
if (handleLogin(req, res)) return;
|
||||
if (req.url === "/api/pages/tree") {
|
||||
sendJson(res, 200, { success: true, data: { items: NODES } });
|
||||
return;
|
||||
}
|
||||
if (req.url === "/api/comments") {
|
||||
const body = JSON.parse(raw || "{}");
|
||||
inFlight++;
|
||||
maxInFlight = Math.max(maxInFlight, inFlight);
|
||||
// Hold the response briefly so concurrent fetches actually overlap.
|
||||
setTimeout(() => {
|
||||
inFlight--;
|
||||
// Every page carries one fresh comment so ordering is observable.
|
||||
sendJson(res, 200, {
|
||||
success: true,
|
||||
data: {
|
||||
items: [
|
||||
{ id: `c-${body.pageId}`, createdAt: "2030-01-01T00:00:00.000Z", content: null },
|
||||
],
|
||||
meta: { nextCursor: null },
|
||||
},
|
||||
});
|
||||
}, 25);
|
||||
return;
|
||||
}
|
||||
sendJson(res, 404, {});
|
||||
});
|
||||
|
||||
const client = new DocmostClient(baseURL, "user@example.com", "pw");
|
||||
const result = await client.checkNewComments(
|
||||
"space-1",
|
||||
"2020-01-01T00:00:00.000Z",
|
||||
"parent",
|
||||
);
|
||||
|
||||
// 13 pages (parent + 12 kids) were scanned; each had a fresh comment.
|
||||
assert.equal(result.checkedPages, 13, "all pages scanned");
|
||||
assert.equal(result.totalNewComments, 13, "one fresh comment per page");
|
||||
// Parallelism: more than one request was in flight at once, but never above the
|
||||
// cap (6). A serial implementation would show maxInFlight === 1.
|
||||
assert.ok(maxInFlight > 1, `expected concurrent fetches, saw max ${maxInFlight}`);
|
||||
assert.ok(maxInFlight <= 6, `concurrency must be bounded, saw ${maxInFlight}`);
|
||||
// Deterministic order: results follow the page-enumeration order (parent first).
|
||||
assert.equal(result.comments[0].pageId, "parent", "results preserve page order");
|
||||
assert.deepEqual(
|
||||
result.comments.map((r) => r.pageId),
|
||||
["parent", ...Array.from({ length: 12 }, (_, i) => `k${i}`)],
|
||||
"result order matches the enumeration order regardless of finish order",
|
||||
);
|
||||
});
|
||||
|
||||
@@ -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