Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 6f3237bc41 | |||
| 8125b7e700 |
@@ -370,10 +370,12 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
);
|
||||
});
|
||||
|
||||
// #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 () => {
|
||||
// #490/#520 reactive branch: a provider CONTEXT-OVERFLOW 400 in onError is
|
||||
// classified, records a distinguishable cause, and stamps the consecutive-overflow
|
||||
// COUNTER (metadata.replayOverflowCount) so the NEXT turn's budgeter trims with
|
||||
// escalating aggression (the recovery that un-bricks the chat). This is a fresh
|
||||
// chat (empty history -> prior streak 0), so the first overflow stamps count 1.
|
||||
it('#490/#520: a context-overflow 400 stamps replayOverflowCount=1 on the finalized row', async () => {
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'error')
|
||||
.mockImplementation(() => undefined as never);
|
||||
@@ -397,11 +399,14 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
metadata: Record<string, unknown>;
|
||||
};
|
||||
expect(patch.status).toBe('error');
|
||||
expect(patch.metadata.replayOverflow).toBe(true);
|
||||
// First overflow on a fresh chat -> k = prior(0) + 1 = 1.
|
||||
expect(patch.metadata.replayOverflowCount).toBe(1);
|
||||
// The legacy boolean is no longer written (the counter supersedes it).
|
||||
expect('replayOverflow' in patch.metadata).toBe(false);
|
||||
expect(patch.metadata.error).toContain('контекстное окно');
|
||||
});
|
||||
|
||||
it('#490: a non-overflow error does NOT stamp replayOverflow', async () => {
|
||||
it('#490/#520: a non-overflow error does NOT stamp the overflow counter', async () => {
|
||||
jest
|
||||
.spyOn(Logger.prototype, 'error')
|
||||
.mockImplementation(() => undefined as never);
|
||||
@@ -412,6 +417,7 @@ describe('AiChatService.stream — abortSignal wiring (#184 F3)', () => {
|
||||
status: string;
|
||||
metadata: Record<string, unknown>;
|
||||
};
|
||||
expect('replayOverflowCount' in patch.metadata).toBe(false);
|
||||
expect('replayOverflow' in patch.metadata).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -30,13 +30,16 @@ import {
|
||||
STEP_LIMIT_NO_ANSWER_MARKER,
|
||||
OUTPUT_DEGENERATION_ERROR,
|
||||
lastAssistantContextTokens,
|
||||
lastAssistantReplayOverflow,
|
||||
lastAssistantReplayOverflowCount,
|
||||
seedActivatedTools,
|
||||
} from './ai-chat.service';
|
||||
import type { AiChatMessage, Workspace } from '@docmost/db/types/entity.types';
|
||||
import { buildSystemPrompt } from './ai-chat.prompt';
|
||||
import type { McpClientsService } from './external-mcp/mcp-clients.service';
|
||||
import { resolveEffectiveReplayThreshold } from './history-budget';
|
||||
import {
|
||||
resolveEffectiveReplayThreshold,
|
||||
REPLAY_MIN_FLOOR_TOKENS,
|
||||
} from './history-budget';
|
||||
|
||||
/**
|
||||
* Unit tests for compactToolOutput: the pure helper that shrinks tool outputs
|
||||
@@ -554,49 +557,123 @@ describe('seedActivatedTools', () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe('lastAssistantReplayOverflow', () => {
|
||||
describe('lastAssistantReplayOverflowCount', () => {
|
||||
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', () => {
|
||||
it('reads the consecutive-overflow count from the LAST assistant turn', () => {
|
||||
expect(
|
||||
lastAssistantReplayOverflow([
|
||||
row('assistant', { replayOverflow: true }),
|
||||
lastAssistantReplayOverflowCount([
|
||||
row('assistant', { replayOverflowCount: 3 }),
|
||||
row('user', null),
|
||||
]),
|
||||
).toBe(true);
|
||||
// A recovered (later, non-overflow) assistant turn clears it.
|
||||
).toBe(3);
|
||||
// A recovered (later, non-overflow) assistant turn resets it to 0 — the read
|
||||
// stops at the most recent assistant row, which carries no count.
|
||||
expect(
|
||||
lastAssistantReplayOverflow([
|
||||
row('assistant', { replayOverflow: true }),
|
||||
lastAssistantReplayOverflowCount([
|
||||
row('assistant', { replayOverflowCount: 3 }),
|
||||
row('user', null),
|
||||
row('assistant', { contextTokens: 5 }),
|
||||
]),
|
||||
).toBe(false);
|
||||
expect(lastAssistantReplayOverflow([])).toBe(false);
|
||||
).toBe(0);
|
||||
expect(lastAssistantReplayOverflowCount([])).toBe(0);
|
||||
});
|
||||
|
||||
// #490 reactive recovery: a prior turn stamped `replayOverflow` must make the
|
||||
// NEXT turn's effective budget the AGGRESSIVE 0.5x cut — that harder trim is
|
||||
// what un-bricks a chat that just 400'd on the context window. This exercises
|
||||
// the exact wiring the service uses: read the stamp, then scale the threshold.
|
||||
it('#490: a prior replayOverflow drives the next turn to the 0.5x aggressive budget', () => {
|
||||
// BACK-COMPAT (#520): an in-flight row written by the pre-#520 boolean stamp
|
||||
// (`replayOverflow: true`, no count) reads as k=1 — the old single 0.5× behavior —
|
||||
// so a chat mid-recovery across the deploy does not regress.
|
||||
it('#520 back-compat: a legacy boolean replayOverflow reads as k=1', () => {
|
||||
expect(
|
||||
lastAssistantReplayOverflowCount([
|
||||
row('assistant', { replayOverflow: true }),
|
||||
row('user', null),
|
||||
]),
|
||||
).toBe(1);
|
||||
// A legacy row with the flag absent/false is k=0.
|
||||
expect(
|
||||
lastAssistantReplayOverflowCount([row('assistant', { contextTokens: 5 })]),
|
||||
).toBe(0);
|
||||
});
|
||||
|
||||
// A corrupt/negative persisted count never yields a negative k.
|
||||
it('clamps a corrupt negative count to 0', () => {
|
||||
expect(
|
||||
lastAssistantReplayOverflowCount([
|
||||
row('assistant', { replayOverflowCount: -4 }),
|
||||
]),
|
||||
).toBe(0);
|
||||
});
|
||||
|
||||
// #490/#520 reactive recovery: the prior consecutive-overflow count `k` drives
|
||||
// the next turn's effective budget to an ESCALATING cut (0.5**k) — each further
|
||||
// consecutive 400 tightens it, which is what un-bricks a chat that keeps
|
||||
// overflowing. This exercises the exact wiring the service uses: read the count,
|
||||
// then scale the threshold.
|
||||
it('#490/#520: the prior count drives the next turn to the escalating aggressive budget', () => {
|
||||
const history = [
|
||||
row('assistant', { replayOverflow: true }),
|
||||
row('assistant', { replayOverflowCount: 1 }),
|
||||
row('user', null),
|
||||
];
|
||||
const priorOverflowed = lastAssistantReplayOverflow(history);
|
||||
expect(priorOverflowed).toBe(true);
|
||||
// Base budget 100k -> aggressive recovery halves it to 50k this turn.
|
||||
expect(resolveEffectiveReplayThreshold(100_000, priorOverflowed)).toBe(50_000);
|
||||
const k = lastAssistantReplayOverflowCount(history);
|
||||
expect(k).toBe(1);
|
||||
// Base budget 100k -> first-overflow recovery halves it to 50k this turn.
|
||||
expect(resolveEffectiveReplayThreshold(100_000, k)).toBe(50_000);
|
||||
// A second consecutive overflow (k=2) quarters it.
|
||||
expect(resolveEffectiveReplayThreshold(100_000, 2)).toBe(25_000);
|
||||
// Odd base floors, not rounds.
|
||||
expect(resolveEffectiveReplayThreshold(99_999, true)).toBe(49_999);
|
||||
// No prior overflow -> the base budget is used verbatim (no aggressive cut).
|
||||
expect(resolveEffectiveReplayThreshold(100_000, false)).toBe(100_000);
|
||||
expect(resolveEffectiveReplayThreshold(99_999, 1)).toBe(49_999);
|
||||
// No prior overflow (k=0) -> the base budget is used verbatim (no cut).
|
||||
expect(resolveEffectiveReplayThreshold(100_000, 0)).toBe(100_000);
|
||||
// An explicit off-switch (null) is never overridden, even on recovery.
|
||||
expect(resolveEffectiveReplayThreshold(null, true)).toBeNull();
|
||||
expect(resolveEffectiveReplayThreshold(null, 3)).toBeNull();
|
||||
});
|
||||
|
||||
// #520 escalation table + convergence: the cut deepens each consecutive overflow
|
||||
// and is CLAMPED at the floor so it converges (un-bricks even against a small
|
||||
// real window), instead of the old fixed single 0.5× that stuck at 50k forever.
|
||||
it('#520: escalates and converges to the floor, un-bricking a small real window', () => {
|
||||
const base = 100_000;
|
||||
expect(resolveEffectiveReplayThreshold(base, 0)).toBe(base);
|
||||
expect(resolveEffectiveReplayThreshold(base, 1)).toBe(50_000);
|
||||
expect(resolveEffectiveReplayThreshold(base, 2)).toBe(25_000);
|
||||
expect(resolveEffectiveReplayThreshold(base, 3)).toBe(12_500);
|
||||
|
||||
// Residual-brick regression (#520): with the flat-default base (100k) and a real
|
||||
// model window of ~40k, the OLD fixed 0.5× stuck at 50k forever (> 40k -> 400s
|
||||
// again, never recovers). The escalating cut drops BELOW 40k after enough
|
||||
// consecutive overflows -> the history finally fits -> the chat un-bricks.
|
||||
const realWindow = 40_000;
|
||||
// k=1 (50k) still exceeds the window — the old behavior's terminal state.
|
||||
expect(resolveEffectiveReplayThreshold(base, 1)).toBeGreaterThan(realWindow);
|
||||
// But escalation converges under the window within a couple more turns.
|
||||
const converged = [2, 3, 4, 5].some(
|
||||
(k) => (resolveEffectiveReplayThreshold(base, k) as number) < realWindow,
|
||||
);
|
||||
expect(converged).toBe(true);
|
||||
|
||||
// Convergence is bounded BELOW by the floor: a large k never trims below it.
|
||||
for (const k of [4, 8, 20, 100]) {
|
||||
expect(resolveEffectiveReplayThreshold(base, k)).toBe(REPLAY_MIN_FLOOR_TOKENS);
|
||||
expect(
|
||||
resolveEffectiveReplayThreshold(base, k) as number,
|
||||
).toBeGreaterThanOrEqual(REPLAY_MIN_FLOOR_TOKENS);
|
||||
}
|
||||
});
|
||||
|
||||
// The floor never RAISES a legitimately small configured budget above itself —
|
||||
// that would re-overflow the very window it was configured for.
|
||||
it('#520: never inflates a small configured budget above itself', () => {
|
||||
const small = 5_000; // below the floor
|
||||
expect(resolveEffectiveReplayThreshold(small, 0)).toBe(small);
|
||||
// Even under escalation the effective threshold never exceeds the base.
|
||||
for (const k of [1, 2, 3, 10]) {
|
||||
expect(
|
||||
resolveEffectiveReplayThreshold(small, k) as number,
|
||||
).toBeLessThanOrEqual(small);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
@@ -930,21 +1007,29 @@ 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', () => {
|
||||
// #490/#520 observability: the replay budgeter's decision is stamped on the turn,
|
||||
// now including the consecutive-overflow COUNTER (#520) the next turn escalates on.
|
||||
it('records replayTrimmedToTokens + replayOverflowCount when provided', () => {
|
||||
const f = flushAssistant([], '', 'error', {
|
||||
error: 'ctx',
|
||||
replayTrimmedToTokens: 42_000,
|
||||
replayOverflow: true,
|
||||
replayOverflowCount: 2,
|
||||
});
|
||||
expect(f.metadata.replayTrimmedToTokens).toBe(42_000);
|
||||
expect(f.metadata.replayOverflow).toBe(true);
|
||||
expect(f.metadata.replayOverflowCount).toBe(2);
|
||||
});
|
||||
|
||||
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);
|
||||
expect('replayOverflowCount' in f.metadata).toBe(false);
|
||||
});
|
||||
|
||||
// A clean finalize (no overflow -> count 0/omitted) leaves NO counter, which the
|
||||
// next turn reads as k=0 — the reset that ends a recovery streak.
|
||||
it('omits replayOverflowCount for a zero/absent count (reset semantics)', () => {
|
||||
const zero = flushAssistant([], '', 'completed', { replayOverflowCount: 0 });
|
||||
expect('replayOverflowCount' in zero.metadata).toBe(false);
|
||||
});
|
||||
|
||||
// #274 observability: the page-change diff the agent saw this turn is persisted
|
||||
|
||||
@@ -140,9 +140,10 @@ const OUTPUT_DEGENERATION_ERROR =
|
||||
|
||||
// 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).
|
||||
// row is ALSO stamped `metadata.replayOverflowCount` (the consecutive-overflow
|
||||
// counter, #520) so the NEXT turn's budgeter trims with escalating aggression (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 =
|
||||
'Диалог превысил контекстное окно модели; история будет агрессивно ' +
|
||||
'сокращена на следующем ходу.';
|
||||
@@ -1192,19 +1193,23 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
}
|
||||
// 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);
|
||||
// Reactive recovery (#490/#520): `k` = how many CONSECUTIVE preceding turns
|
||||
// were rejected for context overflow (stamped by onError). Each consecutive
|
||||
// overflow trims MORE aggressively (resolveEffectiveReplayThreshold scales the
|
||||
// budget by 0.5**k, clamped at the floor) so recovery ESCALATES until the
|
||||
// history fits — the overflowing turn produced no usage signal, so a single
|
||||
// fixed cut may not shrink enough when the real model window is small. This is
|
||||
// what un-bricks a chat that keeps 400'ing on the context window.
|
||||
const priorOverflowCount = lastAssistantReplayOverflowCount(oldHistory);
|
||||
const effectiveThreshold = resolveEffectiveReplayThreshold(
|
||||
replayBudget.thresholdTokens,
|
||||
priorOverflowed,
|
||||
priorOverflowCount,
|
||||
);
|
||||
if (priorOverflowed) {
|
||||
if (priorOverflowCount > 0) {
|
||||
this.logger.warn(
|
||||
`AI chat (chat ${chatId}): previous turn hit context overflow; ` +
|
||||
`applying aggressive replay budget (${effectiveThreshold} tokens).`,
|
||||
`AI chat (chat ${chatId}): ${priorOverflowCount} consecutive context ` +
|
||||
`overflow(s); applying escalated aggressive replay budget ` +
|
||||
`(${effectiveThreshold} tokens).`,
|
||||
);
|
||||
}
|
||||
const preTrim = trimHistoryForReplay(
|
||||
@@ -1212,7 +1217,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
effectiveThreshold,
|
||||
// A prior OVERFLOW means the provider count is stale/absent — force the
|
||||
// char-estimate path by ignoring priorContextTokens on recovery.
|
||||
priorOverflowed ? undefined : priorContextTokens,
|
||||
priorOverflowCount > 0 ? undefined : priorContextTokens,
|
||||
);
|
||||
messages = preTrim.messages;
|
||||
// Observability (#490): record the budgeter's decision on the turn so the UI
|
||||
@@ -1904,7 +1909,12 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy {
|
||||
pageChanged,
|
||||
partsCache,
|
||||
replayTrimmedToTokens,
|
||||
replayOverflow: overflow || undefined,
|
||||
// #520: escalate the consecutive-overflow counter so the NEXT turn
|
||||
// trims MORE aggressively (0.5**k). k grows by 1 each consecutive
|
||||
// overflow; a clean finalize omits the field, resetting it to 0.
|
||||
replayOverflowCount: overflow
|
||||
? priorOverflowCount + 1
|
||||
: undefined,
|
||||
}),
|
||||
);
|
||||
// #184: settle the RUN as failed, carrying the provider/transport cause.
|
||||
@@ -2336,22 +2346,39 @@ export function seedActivatedTools(
|
||||
}
|
||||
|
||||
/**
|
||||
* 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.
|
||||
* How many CONSECUTIVE recent turns were rejected for CONTEXT OVERFLOW (#490/#520):
|
||||
* `k`, read from the most recent assistant row's `metadata.replayOverflowCount`
|
||||
* (stamped by the stream's onError, incremented each consecutive overflow and reset
|
||||
* to 0 on any clean finalize). The next turn's budgeter feeds this to
|
||||
* {@link resolveEffectiveReplayThreshold} to trim with ESCALATING aggression — the
|
||||
* reactive recovery. Only the LAST assistant turn matters (its count already carries
|
||||
* the consecutive streak; an older overflow followed by a clean turn was recovered),
|
||||
* so we stop at the first assistant row scanning backwards.
|
||||
*
|
||||
* BACK-COMPAT: a row written by the pre-#520 boolean stamp (`replayOverflow: true`,
|
||||
* no count) is read as k=1 — the old single 0.5× behavior — so in-flight chats do
|
||||
* not regress across the deploy.
|
||||
*/
|
||||
export function lastAssistantReplayOverflow(
|
||||
export function lastAssistantReplayOverflowCount(
|
||||
history: ReadonlyArray<AiChatMessage>,
|
||||
): boolean {
|
||||
): number {
|
||||
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;
|
||||
const meta = (row.metadata ?? {}) as {
|
||||
replayOverflowCount?: unknown;
|
||||
replayOverflow?: unknown;
|
||||
};
|
||||
if (typeof meta.replayOverflowCount === 'number') {
|
||||
// Guard against a corrupt/negative persisted value.
|
||||
return meta.replayOverflowCount > 0
|
||||
? Math.floor(meta.replayOverflowCount)
|
||||
: 0;
|
||||
}
|
||||
// Back-compat: legacy boolean stamp -> one overflow (0.5× cut).
|
||||
return meta.replayOverflow === true ? 1 : 0;
|
||||
}
|
||||
return false;
|
||||
return 0;
|
||||
}
|
||||
|
||||
/** The last message with role 'user' from a useChat payload, if any. */
|
||||
@@ -2891,9 +2918,11 @@ export function flushAssistant(
|
||||
// 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;
|
||||
// #490/#520 reactive branch: the consecutive context-overflow count for THIS
|
||||
// turn (prior streak + 1) when the provider rejected it for context overflow.
|
||||
// Stamped into metadata so the NEXT turn's budgeter trims with escalating
|
||||
// aggression (0.5**k). Omitted (undefined) on a clean turn, which resets k to 0.
|
||||
replayOverflowCount?: number;
|
||||
},
|
||||
): AssistantFlush {
|
||||
const finished = capturedSteps ?? [];
|
||||
@@ -2946,7 +2975,8 @@ export function flushAssistant(
|
||||
metadata.maxContextTokens = extra.maxContextTokens;
|
||||
if (extra?.replayTrimmedToTokens)
|
||||
metadata.replayTrimmedToTokens = extra.replayTrimmedToTokens;
|
||||
if (extra?.replayOverflow) metadata.replayOverflow = true;
|
||||
if (extra?.replayOverflowCount && extra.replayOverflowCount > 0)
|
||||
metadata.replayOverflowCount = extra.replayOverflowCount;
|
||||
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
|
||||
|
||||
@@ -1,6 +1,11 @@
|
||||
import { randomBytes } from 'crypto';
|
||||
import { Client } from 'pg';
|
||||
import { flushAssistant, serializeSteps } from './ai-chat.service';
|
||||
import {
|
||||
flushAssistant,
|
||||
serializeSteps,
|
||||
lastAssistantReplayOverflowCount,
|
||||
} from './ai-chat.service';
|
||||
import type { AiChatMessage } from '@docmost/db/types/entity.types';
|
||||
|
||||
/**
|
||||
* #490 write-volume regression — an OBSERVABLE-PROPERTY test on a LIVE Postgres,
|
||||
@@ -207,3 +212,111 @@ describe('#490 write-volume on a live Postgres (pg_current_wal_lsn delta)', () =
|
||||
expect(v2).toBeLessThan(v1 * 0.75);
|
||||
}, 120_000);
|
||||
});
|
||||
|
||||
/**
|
||||
* #520 reactive-recovery COUNTER lifecycle on a LIVE Postgres — proves the
|
||||
* consecutive-overflow count survives a real jsonb metadata round-trip (the persist
|
||||
* path), not just an in-memory object. flushAssistant BUILDS the row metadata, we
|
||||
* WRITE it to a jsonb column, READ it back, then reconstruct the assistant row and
|
||||
* run lastAssistantReplayOverflowCount over it — exactly the read the next turn does.
|
||||
*
|
||||
* The lifecycle proven end-to-end through pg:
|
||||
* - consecutive overflows INCREMENT k (1 -> 2 -> 3);
|
||||
* - a CLEAN finalize omits the field, which the reader treats as a RESET to 0;
|
||||
* - a legacy boolean row (`replayOverflow: true`) reads back as k=1 (back-compat).
|
||||
*/
|
||||
describe('#520 overflow-counter lifecycle on a live Postgres (jsonb round-trip)', () => {
|
||||
let client: Client | undefined;
|
||||
let available = false;
|
||||
|
||||
beforeAll(async () => {
|
||||
try {
|
||||
client = new Client(CONN);
|
||||
await client.connect();
|
||||
await client.query('SELECT 1');
|
||||
available = true;
|
||||
} catch {
|
||||
available = false;
|
||||
client = undefined;
|
||||
}
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await client?.end().catch(() => undefined);
|
||||
});
|
||||
|
||||
// Round-trip an arbitrary metadata object through a real jsonb column and read it
|
||||
// back as the reconstructed assistant row the next turn would load.
|
||||
async function roundTrip(
|
||||
c: Client,
|
||||
metadata: unknown,
|
||||
): Promise<AiChatMessage> {
|
||||
await c.query('UPDATE _wal_counter SET metadata=$1 WHERE id=1', [
|
||||
JSON.stringify(metadata),
|
||||
]);
|
||||
const back = (await c.query('SELECT metadata FROM _wal_counter WHERE id=1'))
|
||||
.rows[0].metadata as Record<string, unknown>;
|
||||
return { role: 'assistant', metadata: back } as unknown as AiChatMessage;
|
||||
}
|
||||
|
||||
it('increments across consecutive overflows, resets on a clean turn, and honors the legacy boolean', async () => {
|
||||
if (!available || !client) {
|
||||
console.warn('SKIP: gitmost-test-pg not reachable; skipping counter test.');
|
||||
return;
|
||||
}
|
||||
const c = client;
|
||||
await c.query('DROP TABLE IF EXISTS _wal_counter');
|
||||
await c.query('CREATE TABLE _wal_counter(id int primary key, metadata jsonb)');
|
||||
await c.query("INSERT INTO _wal_counter VALUES (1, '{}'::jsonb)");
|
||||
|
||||
// Turn 1 overflow: prior streak 0 -> stamp k=1 (as the service does: prior+1).
|
||||
let prior = lastAssistantReplayOverflowCount([]); // fresh chat
|
||||
expect(prior).toBe(0);
|
||||
let row = await roundTrip(
|
||||
c,
|
||||
flushAssistant([], '', 'error', {
|
||||
error: 'ctx',
|
||||
replayOverflowCount: prior + 1,
|
||||
}).metadata,
|
||||
);
|
||||
prior = lastAssistantReplayOverflowCount([row]);
|
||||
expect(prior).toBe(1);
|
||||
|
||||
// Turn 2 overflow: prior 1 -> stamp k=2.
|
||||
row = await roundTrip(
|
||||
c,
|
||||
flushAssistant([], '', 'error', {
|
||||
error: 'ctx',
|
||||
replayOverflowCount: prior + 1,
|
||||
}).metadata,
|
||||
);
|
||||
prior = lastAssistantReplayOverflowCount([row]);
|
||||
expect(prior).toBe(2);
|
||||
|
||||
// Turn 3 overflow: prior 2 -> stamp k=3.
|
||||
row = await roundTrip(
|
||||
c,
|
||||
flushAssistant([], '', 'error', {
|
||||
error: 'ctx',
|
||||
replayOverflowCount: prior + 1,
|
||||
}).metadata,
|
||||
);
|
||||
prior = lastAssistantReplayOverflowCount([row]);
|
||||
expect(prior).toBe(3);
|
||||
|
||||
// Turn 4 CLEAN finalize: no overflow -> the field is omitted -> reset to 0.
|
||||
row = await roundTrip(
|
||||
c,
|
||||
flushAssistant([], 'all good', 'completed', { finishReason: 'stop' })
|
||||
.metadata,
|
||||
);
|
||||
expect('replayOverflowCount' in (row.metadata as object)).toBe(false);
|
||||
expect(lastAssistantReplayOverflowCount([row])).toBe(0);
|
||||
|
||||
// Back-compat: a row persisted by the pre-#520 boolean stamp reads back as k=1.
|
||||
row = await roundTrip(c, { replayOverflow: true });
|
||||
expect(lastAssistantReplayOverflowCount([row])).toBe(1);
|
||||
|
||||
await c.query('DROP TABLE IF EXISTS _wal_counter');
|
||||
}, 60_000);
|
||||
});
|
||||
|
||||
@@ -1,14 +1,79 @@
|
||||
import type { ModelMessage } from 'ai';
|
||||
import {
|
||||
resolveReplayBudget,
|
||||
resolveEffectiveReplayThreshold,
|
||||
isContextOverflowError,
|
||||
estimateMessagesTokens,
|
||||
trimHistoryForReplay,
|
||||
REPLAY_BUDGET_DEFAULT_TOKENS,
|
||||
REPLAY_MIN_FLOOR_TOKENS,
|
||||
REPLAY_TRUNCATION_MARKER,
|
||||
REPLAY_TURN_COLLAPSED_MARKER,
|
||||
} from './history-budget';
|
||||
|
||||
describe('resolveEffectiveReplayThreshold (#520 iterative escalation)', () => {
|
||||
// The escalation table: each consecutive overflow (k) deepens the cut by 0.5×.
|
||||
it('scales the base by 0.5**k, flooring (not rounding) fractional tokens', () => {
|
||||
const base = 100_000;
|
||||
expect(resolveEffectiveReplayThreshold(base, 0)).toBe(base); // k=0: unchanged
|
||||
expect(resolveEffectiveReplayThreshold(base, 1)).toBe(50_000); // 0.5×
|
||||
expect(resolveEffectiveReplayThreshold(base, 2)).toBe(25_000); // 0.25×
|
||||
expect(resolveEffectiveReplayThreshold(base, 3)).toBe(12_500); // 0.125×
|
||||
// Floors, not rounds.
|
||||
expect(resolveEffectiveReplayThreshold(99_999, 1)).toBe(49_999);
|
||||
});
|
||||
|
||||
it('passes a null base (trimming OFF) through unchanged for any k', () => {
|
||||
for (const k of [0, 1, 2, 5, 100]) {
|
||||
expect(resolveEffectiveReplayThreshold(null, k)).toBeNull();
|
||||
}
|
||||
});
|
||||
|
||||
// The crux of #520: convergence. A large k is clamped at REPLAY_MIN_FLOOR_TOKENS,
|
||||
// so the escalation CONVERGES to a small-but-usable budget instead of trimming to
|
||||
// zero — and, unlike the old fixed 0.5× that stuck at 50k, it drops far enough to
|
||||
// fit a small real model window.
|
||||
it('clamps a large k at the floor (converges, never below)', () => {
|
||||
const base = 100_000;
|
||||
for (const k of [4, 6, 10, 50, 200]) {
|
||||
const t = resolveEffectiveReplayThreshold(base, k) as number;
|
||||
expect(t).toBe(REPLAY_MIN_FLOOR_TOKENS);
|
||||
expect(t).toBeGreaterThanOrEqual(REPLAY_MIN_FLOOR_TOKENS);
|
||||
}
|
||||
});
|
||||
|
||||
// Residual-brick regression (#520): flat-default base 100k, real window ~40k. The
|
||||
// OLD fixed single 0.5× stuck at 50k > 40k forever (re-overflows every turn — the
|
||||
// brick). The iterative cut drops BELOW 40k after a couple more consecutive
|
||||
// overflows, so the history finally fits and the chat un-bricks.
|
||||
it('un-bricks: escalation drops below a small real window the fixed 0.5× never could', () => {
|
||||
const base = 100_000;
|
||||
const realWindow = 40_000;
|
||||
// The old terminal state: 0.5× = 50k, still above the window.
|
||||
expect(resolveEffectiveReplayThreshold(base, 1)).toBeGreaterThan(realWindow);
|
||||
// Escalation converges under the window.
|
||||
const converged = [2, 3, 4, 5].some(
|
||||
(k) => (resolveEffectiveReplayThreshold(base, k) as number) < realWindow,
|
||||
);
|
||||
expect(converged).toBe(true);
|
||||
// MUTATION SENTINEL: reverting `** k` to `** 1` (fixed 0.5×) makes every k yield
|
||||
// 50k, so `converged` above would be FALSE and this test reddens. Removing the
|
||||
// floor reddens the clamp test instead.
|
||||
});
|
||||
|
||||
// The floor never RAISES a legitimately small configured budget above itself
|
||||
// (min(floor, base)); doing so would re-overflow the very small window it was set
|
||||
// for. So a base BELOW the floor is passed through unchanged and never inflated.
|
||||
it('never inflates a small configured budget above itself', () => {
|
||||
const small = 5_000; // below REPLAY_MIN_FLOOR_TOKENS
|
||||
expect(resolveEffectiveReplayThreshold(small, 0)).toBe(small);
|
||||
for (const k of [1, 2, 3, 10]) {
|
||||
const t = resolveEffectiveReplayThreshold(small, k) as number;
|
||||
expect(t).toBeLessThanOrEqual(small);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('resolveReplayBudget', () => {
|
||||
it('uses floor(0.7 x window) for a configured window (no cap)', () => {
|
||||
// 0.7 x 60k = 42k
|
||||
|
||||
@@ -22,10 +22,25 @@ 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.
|
||||
* Per-step fraction of the normal budget applied on the REACTIVE re-trim after a
|
||||
* provider context-overflow 400 — the preventive estimate under-counted, so cut
|
||||
* harder. This is now applied ITERATIVELY: with `k` consecutive overflow turns the
|
||||
* budget is scaled by `fraction ** k` (k=1 -> 0.5×, k=2 -> 0.25×, …), so recovery
|
||||
* ESCALATES turn over turn until the replayed history finally fits, instead of the
|
||||
* old single fixed 0.5× cut that could never un-brick a chat whose real model
|
||||
* window is smaller than 0.5 × the (unconfigured, flat-default) base budget (#520).
|
||||
*/
|
||||
export const REPLAY_AGGRESSIVE_FRACTION = 0.5;
|
||||
/**
|
||||
* Lower bound (tokens) on the escalating reactive budget: the iterative cut is
|
||||
* clamped here so it CONVERGES (a fixed floor, not an ever-shrinking value that
|
||||
* would eventually trim everything). Rationale: below ~8k tokens a chat cannot
|
||||
* carry meaningful recent context, and even a small real model window comfortably
|
||||
* fits this much — keep-recent-turns still applies on top, so a handful of recent
|
||||
* turns survive. It is never applied so as to RAISE a legitimately small configured
|
||||
* budget (that would re-overflow a tiny window); see resolveEffectiveReplayThreshold.
|
||||
*/
|
||||
export const REPLAY_MIN_FLOOR_TOKENS = 8_000;
|
||||
/**
|
||||
* Turns (a user message + its assistant/tool replies) kept FULL at the tail,
|
||||
* including the current one — never trimmed. Older turns are compacted first.
|
||||
@@ -85,22 +100,34 @@ export function resolveReplayBudget(rawContextWindow: unknown): ReplayBudget {
|
||||
}
|
||||
|
||||
/**
|
||||
* The effective replay threshold for THIS turn, given the base budget and whether
|
||||
* the PREVIOUS turn hit a context-overflow 400 (the reactive-recovery signal,
|
||||
* `metadata.replayOverflow`). On recovery the base budget is scaled down by
|
||||
* {@link REPLAY_AGGRESSIVE_FRACTION}: the overflowing turn produced no usage
|
||||
* signal, so the preventive estimate under-counted and a normal-threshold trim may
|
||||
* not shrink enough to fit — this harder cut is what un-bricks the chat.
|
||||
* The effective replay threshold for THIS turn, given the base budget and `k` — the
|
||||
* number of CONSECUTIVE preceding turns that hit a context-overflow 400 (the
|
||||
* reactive-recovery signal `metadata.replayOverflowCount`, read from the last
|
||||
* assistant row). On recovery the base budget is scaled down ITERATIVELY by
|
||||
* {@link REPLAY_AGGRESSIVE_FRACTION} ** k and clamped at {@link REPLAY_MIN_FLOOR_TOKENS}:
|
||||
* - k=0 -> base unchanged (no overflow: nothing to recover from).
|
||||
* - k=1 -> floor(0.5 × base); k=2 -> floor(0.25 × base); … each further consecutive
|
||||
* overflow tightens the cut, so recovery ESCALATES until the history fits.
|
||||
* - the escalation is clamped at the floor so it CONVERGES — this is what un-bricks
|
||||
* a chat whose real model window is smaller than a single 0.5× cut of the base
|
||||
* (e.g. an unconfigured window: flat-default base 100k, real window <50k) (#520).
|
||||
*
|
||||
* The overflowing turn produced no usage signal, so the preventive estimate
|
||||
* under-counted and a normal-threshold (or single fixed 0.5×) trim may not shrink
|
||||
* enough to fit; the escalating cut is what recovers such a chat.
|
||||
*
|
||||
* A `null` base budget (trimming OFF) is passed through unchanged: an explicit
|
||||
* off-switch is never overridden by the recovery path.
|
||||
* off-switch is never overridden by the recovery path. The floor is applied as
|
||||
* `min(floor, base)` so it never RAISES a legitimately small configured budget
|
||||
* above itself (which would re-overflow the same small window it was set for).
|
||||
*/
|
||||
export function resolveEffectiveReplayThreshold(
|
||||
thresholdTokens: number | null,
|
||||
priorOverflowed: boolean,
|
||||
k: number,
|
||||
): number | null {
|
||||
if (!priorOverflowed || thresholdTokens == null) return thresholdTokens;
|
||||
return Math.floor(thresholdTokens * REPLAY_AGGRESSIVE_FRACTION);
|
||||
if (thresholdTokens == null || k <= 0) return thresholdTokens;
|
||||
const scaled = Math.floor(thresholdTokens * REPLAY_AGGRESSIVE_FRACTION ** k);
|
||||
return Math.max(scaled, Math.min(REPLAY_MIN_FLOOR_TOKENS, thresholdTokens));
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
Reference in New Issue
Block a user