Skip to content

Commit 9116203

Browse files
yetvalyetval
committed
fix(reply): project preflight compaction gate by next-input on fresh tokens
The fresh-tokens path of runPreflightCompactionIfNeeded fed the prompt-only entry.totalTokens snapshot straight into the budget threshold check, dropping the current user prompt estimate and the previous turn's output. The sibling memory-flush gate and this function's own stale branch already project base + output + estimate via resolveEffectivePromptTokens, so the preflight gate under-triggered and let over-budget requests through to overflow-retry. Project the fresh persisted base the same way: read transcript output when near the threshold (mirroring the memory-flush gate's buffer) and run the fresh base through resolveEffectivePromptTokens before the threshold check.
1 parent b0998f7 commit 9116203

2 files changed

Lines changed: 98 additions & 11 deletions

File tree

‎src/auto-reply/reply/agent-runner-memory.test.ts‎

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1314,6 +1314,79 @@ describe("runMemoryFlushIfNeeded", () => {
13141314
expect(compactCall.authProfileId).toBe("anthropic:[email protected]");
13151315
expect(compactCall.contextTokenBudget).toBe(258_000);
13161316
});
1317+
it("preflight compacts a fresh session when the current prompt estimate pushes the next request over budget", async () => {
1318+
registerMemoryFlushPlanResolverForTest(() => ({
1319+
softThresholdTokens: 0,
1320+
forceFlushTranscriptBytes: 1_000_000_000,
1321+
reserveTokensFloor: 10,
1322+
prompt: "Pre-compaction memory flush.\nNO_REPLY",
1323+
systemPrompt: "Write memory to memory/YYYY-MM-DD.md.",
1324+
relativePath: "memory/2023-11-14.md",
1325+
}));
1326+
const sessionEntry: SessionEntry = {
1327+
sessionId: "session",
1328+
updatedAt: Date.now(),
1329+
totalTokens: 985,
1330+
totalTokensFresh: true,
1331+
compactionCount: 0,
1332+
};
1333+
1334+
await runPreflightCompactionIfNeeded({
1335+
cfg: { agents: { defaults: { compaction: { memoryFlush: {} } } } },
1336+
followupRun: createTestFollowupRun({
1337+
provider: "anthropic",
1338+
model: "claude",
1339+
sessionKey: "agent:main:main",
1340+
}),
1341+
promptForEstimate: "Please summarize the entire design discussion above. ".repeat(8),
1342+
defaultModel: "anthropic/claude",
1343+
agentCfgContextTokens: 1000,
1344+
sessionEntry,
1345+
sessionStore: { "agent:main:main": sessionEntry },
1346+
sessionKey: "agent:main:main",
1347+
isHeartbeat: false,
1348+
replyOperation: createReplyOperation(),
1349+
});
1350+
1351+
expect(compactEmbeddedAgentSessionMock).toHaveBeenCalledTimes(1);
1352+
});
1353+
it("does not preflight compact a fresh session when only accumulated output tokens are large and the latest output keeps the request under budget", async () => {
1354+
registerMemoryFlushPlanResolverForTest(() => ({
1355+
softThresholdTokens: 0,
1356+
forceFlushTranscriptBytes: 1_000_000_000,
1357+
reserveTokensFloor: 10,
1358+
prompt: "Pre-compaction memory flush.\nNO_REPLY",
1359+
systemPrompt: "Write memory to memory/YYYY-MM-DD.md.",
1360+
relativePath: "memory/2023-11-14.md",
1361+
}));
1362+
const sessionEntry: SessionEntry = {
1363+
sessionId: "session",
1364+
updatedAt: Date.now(),
1365+
totalTokens: 985,
1366+
outputTokens: 50_000,
1367+
totalTokensFresh: true,
1368+
compactionCount: 0,
1369+
};
1370+
1371+
await runPreflightCompactionIfNeeded({
1372+
cfg: { agents: { defaults: { compaction: { memoryFlush: {} } } } },
1373+
followupRun: createTestFollowupRun({
1374+
provider: "anthropic",
1375+
model: "claude",
1376+
sessionKey: "agent:main:main",
1377+
}),
1378+
promptForEstimate: "",
1379+
defaultModel: "anthropic/claude",
1380+
agentCfgContextTokens: 1000,
1381+
sessionEntry,
1382+
sessionStore: { "agent:main:main": sessionEntry },
1383+
sessionKey: "agent:main:main",
1384+
isHeartbeat: false,
1385+
replyOperation: createReplyOperation(),
1386+
});
1387+
1388+
expect(compactEmbeddedAgentSessionMock).not.toHaveBeenCalled();
1389+
});
13171390
it("updates the active preflight run after transcript rotation", async () => {
13181391
const sessionFile = path.join(rootDir, "session.jsonl");
13191392
const successorFile = path.join(rootDir, "session-rotated.jsonl");
@@ -2011,7 +2084,7 @@ describe("runMemoryFlushIfNeeded", () => {
20112084
const compactCall = requireCompactEmbeddedAgentSessionCall();
20122085
expect(compactCall.sessionId).toBe("session");
20132086
expect(compactCall.trigger).toBe("budget");
2014-
expect(compactCall.currentTokenCount).toBe(10);
2087+
expect(compactCall.currentTokenCount).toBe(12);
20152088
expect(compactCall.sessionFile).toContain("large-session.jsonl");
20162089
});
20172090

‎src/auto-reply/reply/agent-runner-memory.ts‎

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -756,10 +756,24 @@ export async function runPreflightCompactionIfNeeded(params: {
756756
const promptTokenEstimate = estimatePromptTokensForMemoryFlush(
757757
params.promptForEstimate ?? params.followupRun.prompt,
758758
);
759+
const serverCompactionThreshold = resolveResponsesServerCompactionThreshold({
760+
cfg: params.cfg,
761+
provider: params.followupRun.run.provider,
762+
modelId: params.followupRun.run.model ?? params.defaultModel,
763+
});
764+
const threshold = Math.max(
765+
contextWindowTokens - reserveTokensFloor - softThresholdTokens,
766+
serverCompactionThreshold ?? 0,
767+
);
768+
const freshNeedsOutputRead =
769+
typeof freshPersistedTokens === "number" &&
770+
typeof promptTokenEstimate === "number" &&
771+
threshold > 0 &&
772+
freshPersistedTokens + promptTokenEstimate >= threshold - TRANSCRIPT_OUTPUT_READ_BUFFER_TOKENS;
759773
const maxActiveTranscriptBytes = resolveMaxActiveTranscriptBytes(params.cfg);
760774
const shouldCheckActiveTranscriptBytes = typeof maxActiveTranscriptBytes === "number";
761775
const transcriptUsageTokens =
762-
typeof freshPersistedTokens === "number"
776+
typeof freshPersistedTokens === "number" && !freshNeedsOutputRead
763777
? undefined
764778
: await estimatePromptTokensFromSessionTranscript({
765779
sessionId: entry.sessionId,
@@ -801,24 +815,24 @@ export async function runPreflightCompactionIfNeeded(params: {
801815
promptTokenEstimate,
802816
)
803817
: undefined;
818+
const freshProjectedTokenCount =
819+
typeof freshPersistedTokens === "number"
820+
? resolveEffectivePromptTokens(
821+
freshPersistedTokens,
822+
transcriptOutputTokens,
823+
promptTokenEstimate,
824+
)
825+
: undefined;
804826
const projectedTokenCount = Math.max(
805827
usageProjectedTokenCount ?? 0,
828+
freshProjectedTokenCount ?? 0,
806829
stalePersistedPromptTokens ?? 0,
807830
);
808831
const tokenCountForCompaction =
809832
Number.isFinite(projectedTokenCount) && projectedTokenCount > 0
810833
? projectedTokenCount
811834
: undefined;
812835

813-
const serverCompactionThreshold = resolveResponsesServerCompactionThreshold({
814-
cfg: params.cfg,
815-
provider: params.followupRun.run.provider,
816-
modelId: params.followupRun.run.model ?? params.defaultModel,
817-
});
818-
const threshold = Math.max(
819-
contextWindowTokens - reserveTokensFloor - softThresholdTokens,
820-
serverCompactionThreshold ?? 0,
821-
);
822836
logVerbose(
823837
`preflightCompaction check: sessionKey=${params.sessionKey} ` +
824838
`tokenCount=${tokenCountForCompaction ?? freshPersistedTokens ?? "undefined"} ` +

0 commit comments

Comments
 (0)