From 156f2063dce13623c23341154e55946a54cb92a6 Mon Sep 17 00:00:00 2001 From: Jiachen Yu Date: Sat, 24 Oct 2026 14:25:11 +0800 Subject: [PATCH] fix(runtime): keep the execution budget reminder in history The execution budget reminder was a request-only marker appended after the last user message. The Anthropic adapter places the conversation cache breakpoint on that message, so each reminder request wrote a cache entry that ended at a marker the next request no longer contained. In the final two minutes every request carried such a marker, so cache reads stopped advancing and the uncached tail grew with every step until it fell out of the lookback window and the whole history was written again. runBeforeLLM now accumulates appendMessage decisions instead of ending the decision pipeline at the first append, so the background, session identity and todo reminders and the request filter no longer skip each other. All markers are committed once after the pipeline: at most one durable replacement carrying placed markers, then every tail marker in registration order with a single history flush. A durable replacement after a registered append fails closed. Compaction commit metadata and canonical tail validation accept more than one placed marker through a runner-owned placedMarkerCount (absent means 1, so existing metadata stays readable). The execution budget reminder becomes a hidden durable custom marker appended like the other reminders, so every provider request is an append-only extension of the previous one. Progress is derived from the latest reminder of the same deadline in canonical history: a repeated request does not duplicate it, a compaction that summarized reminders away announces the budget again, and in the final window a reminder is only appended when its 10-second text changes. The model-visible text, thresholds and position are unchanged. --- .../agent-core/src/pi-turn-runner/hooks.ts | 13 +- packages/agent-core/src/pi-turn-runner/llm.ts | 240 ++++++++--- .../history/canonical-history-provider.ts | 30 +- .../compaction/context-compaction.test.ts | 121 +++++- .../compaction/context-compaction.ts | 76 +++- .../agent-host/execution/executor.test.ts | 196 ++++----- .../agent-host/execution/executor.ts | 10 +- .../runner/execution-budget-reminder.test.ts | 403 ++++++++++++++++++ .../runner/execution-budget-reminder.ts | 118 +++-- .../turn-system/production-composition.ts | 7 +- ...on-budget-cache-prefix.integration.test.ts | 144 +++++++ release/public-source.json | 2 + test/byok.test.mjs | 14 +- test/vitest-suites.json | 2 + 14 files changed, 1114 insertions(+), 262 deletions(-) create mode 100644 packages/local-runtime-v2/src/service/turn-system/agent-host/runner/execution-budget-reminder.test.ts create mode 100644 packages/local-runtime-v2/test/integration/execution-budget-cache-prefix.integration.test.ts diff --git a/packages/agent-core/src/pi-turn-runner/hooks.ts b/packages/agent-core/src/pi-turn-runner/hooks.ts index 9702ff9fb..0f6f3cef7 100644 --- a/packages/agent-core/src/pi-turn-runner/hooks.ts +++ b/packages/agent-core/src/pi-turn-runner/hooks.ts @@ -77,6 +77,13 @@ export interface PiBeforeLlmCallReplaceMetadata { replacementSourceIndexes?: number[]; /** Reserved lineage for a placement-owned exact current-user tail. */ currentUserSourceIndex?: number; + /** + * Reserved count of hidden markers placed immediately before the + * placement-owned current-user tail. Maintained by the runner alongside + * `currentUserSourceIndex`; hooks must not set it. Absent means 1 when + * `currentUserSourceIndex` is present, keeping older persisted metadata valid. + */ + placedMarkerCount?: number; } export type PiBeforeLlmCallAppendMessage = Omit & { @@ -109,7 +116,11 @@ export type PiBeforeLlmCallHookDecision = | { /** * Durably appends one hidden typed context message before the Provider request. - * A successful append ends the before-LLM decision pipeline. + * Appends accumulate: later hooks still run and receive the registered + * markers in `messages` and `canonicalMessages`. All markers are committed + * once after the pipeline (placed markers inside the durable replacement, + * tail markers in registration order with one history flush). A durable + * `replaceMessages` after a registered append aborts the request. */ type: 'appendMessage'; message: PiBeforeLlmCallAppendMessage; diff --git a/packages/agent-core/src/pi-turn-runner/llm.ts b/packages/agent-core/src/pi-turn-runner/llm.ts index d9bcfcb8c..77b8a8990 100644 --- a/packages/agent-core/src/pi-turn-runner/llm.ts +++ b/packages/agent-core/src/pi-turn-runner/llm.ts @@ -78,16 +78,10 @@ export function setLLMHook(agent: Agent, turn: turnState, history: turnHistory): if (checkpoint.type === 'requestOnly') { return checkpoint.messages; } - if (checkpoint.type !== 'replace' && checkpoint.type !== 'append') return messages; + if (checkpoint.type !== 'commit') return messages; - const replacement = - checkpoint.type === 'replace' - ? { - messages: checkpoint.durableMessages, - metadata: checkpoint.metadata, - ...(checkpoint.afterCommit ? { afterCommit: checkpoint.afterCommit } : {}), - } - : checkpoint.durableReplacement; + let requestMessages = checkpoint.messages; + const replacement = checkpoint.durableReplacement; if (replacement) { // Pi passes the provider-bound context array into transformContext, but // Agent.state.messages is the transcript later event reducers read. A @@ -108,18 +102,29 @@ export function setLLMHook(agent: Agent, turn: turnState, history: turnHistory): if (invalidAppendReason) { throw new Error(`after_commit append message invalid: ${invalidAppendReason}`); } - const appendMessage = checkpoint.type === 'append' ? checkpoint.messages.at(-1) : undefined; - checkpoint.messages = appendMessage - ? [...checkpoint.messages.slice(0, -1), postCommitMessage, appendMessage] - : [...checkpoint.messages, postCommitMessage]; + // The post-commit message belongs to the durable replacement, so it + // precedes every tail marker registered by later hooks, both in the + // request and in history. A request-only filter may have dropped the + // tail markers from the request; then the message simply goes last. + const tailMarkers = new Set(checkpoint.tailMessages); + const firstTailIndex = requestMessages.findIndex((message) => tailMarkers.has(message)); + requestMessages = + firstTailIndex < 0 + ? [...requestMessages, postCommitMessage] + : [ + ...requestMessages.slice(0, firstTailIndex), + postCommitMessage, + ...requestMessages.slice(firstTailIndex), + ]; messages.push(postCommitMessage); agent.state.messages = [...agent.state.messages, postCommitMessage]; await history.flushTail(); } } - if (checkpoint.type === 'replace') return checkpoint.messages; - + if (checkpoint.tailMessages.length === 0) return requestMessages; + // Every tail marker is committed in registration order with a single + // delivery ack; any failure keeps the provider call fail-closed. messages.push(...checkpoint.tailMessages); agent.state.messages = [...agent.state.messages, ...checkpoint.tailMessages]; try { @@ -128,7 +133,7 @@ export function setLLMHook(agent: Agent, turn: turnState, history: turnHistory): const message = err instanceof Error ? err.message : String(err); throw new Error(`before_llm_append_persistence_failed: ${message}`); } - return checkpoint.messages; + return requestMessages; }; } @@ -141,37 +146,35 @@ async function runBeforeLLM( | { type: 'unchanged'; messages: AgentMessage[] } | { type: 'requestOnly'; messages: AgentMessage[] } | { - type: 'replace'; + /** + * Durable work to commit before the provider call: at most one durable + * replacement (already carrying placed markers) followed by every tail + * marker in hook registration order. `messages` is the final request view. + */ + type: 'commit'; messages: AgentMessage[]; - durableMessages: AgentMessage[]; - metadata: PiBeforeLlmCallReplaceMetadata; - afterCommit?: PiAfterLlmReplacementCommit; - } - | { - type: 'append'; - messages: AgentMessage[]; - message: PiBeforeLlmCallAppendMessage; - tailMessages: AgentMessage[]; - durableReplacement?: { - messages: AgentMessage[]; - metadata: PiBeforeLlmCallReplaceMetadata; - afterCommit?: PiAfterLlmReplacementCommit; - }; + tailMessages: PiBeforeLlmCallAppendMessage[]; + durableReplacement?: DurableReplacement; } | { type: 'respond'; text: string; reason: string } | { type: 'abort'; reason: string } > { const initialCurrentUser = captureInitialCurrentUser(messages, phase); + // Two views flow through the pipeline: + // - `currentMessages` is the Provider request view. Request-only hooks may + // rewrite it; every registered append marker is inserted here at once so + // later hooks (e.g. final-request budget admission) see the real request. + // - `canonicalMessages` is the canonical history as it will look after this + // checkpoint commits: the durable replacement base (already carrying placed + // markers) followed by every pending tail marker in registration order. let currentMessages = [...messages]; let canonicalMessages = [...messages]; let requestMessagesReplaced = false; - let durableReplacement: - | { - messages: AgentMessage[]; - metadata: PiBeforeLlmCallReplaceMetadata; - afterCommit?: PiAfterLlmReplacementCommit; - } - | undefined; + let durableReplacement: DurableReplacement | undefined; + // Successful appends no longer end the pipeline; tail markers accumulate here + // and are committed with one history flush after every hook has run. + const pendingTailMessages: PiBeforeLlmCallAppendMessage[] = []; + let pendingAppendCount = 0; for (const hook of turn.hooks.beforeLLM) { let decision: Awaited>; try { @@ -309,11 +312,14 @@ async function runBeforeLLM( }; } if (decision.placement === 'before-current-user') { + // Pending tail markers stay at the request tail after the real user, the + // same position they take in canonical history once committed. + const pendingTails = new Set(pendingTailMessages); const placement = placeBeforeCurrentUser({ phase, initialMessages: messages, initialCurrentUser, - currentMessages, + currentMessages: currentMessages.filter((message) => !pendingTails.has(message)), durableReplacement, marker: decision.message, }); @@ -331,24 +337,48 @@ async function runBeforeLLM( } if (placement.type === 'defer') continue; if (placement.type === 'placed') { - return { - type: 'append', - messages: placement.messages, - message: decision.message, - tailMessages: [], - durableReplacement: placement.durableReplacement, + currentMessages = [ + ...placement.messages, + ...currentMessages.filter((message) => pendingTails.has(message)), + ]; + durableReplacement = { + ...placement.durableReplacement, + ...(durableReplacement?.afterCommit + ? { afterCommit: durableReplacement.afterCommit } + : {}), }; + canonicalMessages = [...durableReplacement.messages, ...pendingTailMessages]; + pendingAppendCount += 1; + logAppendRegistered(turn, phase, decision.reason, 'before-current-user'); + continue; } // 'tail': a continuation has no current user to precede, so the ordinary // tail append already keeps real user input last. } - return { - type: 'append', - messages: [...currentMessages, decision.message], - message: decision.message, - tailMessages: [decision.message], - ...(durableReplacement ? { durableReplacement } : {}), - }; + pendingTailMessages.push(decision.message); + currentMessages = [...currentMessages, decision.message]; + canonicalMessages = [...canonicalMessages, decision.message]; + pendingAppendCount += 1; + logAppendRegistered(turn, phase, decision.reason, 'tail'); + continue; + } + if (pendingAppendCount > 0) { + // A durable replacement after a registered marker would have to rebase + // markers the hook never saw onto a history it rewrote. Production + // registers context compaction before every appending reminder, so this + // order is a composition bug; fail closed instead of guessing a rebase. + const reason = + 'beforeLlmCall durable replacement after a pending appendMessage is unsupported'; + turn.logger.error( + { + session_id: turn.input.sessionId, + turn_id: turn.input.turnId, + phase, + reason, + }, + '[pi-turn-runner] beforeLlmCall replacement rejected', + ); + return { type: 'abort', reason }; } const invalidReplacementReason = validReplacement(decision.metadata); if (invalidReplacementReason) { @@ -396,13 +426,12 @@ async function runBeforeLLM( '[pi-turn-runner] beforeLlmCall replaced context messages', ); } - if (durableReplacement) { + if (durableReplacement || pendingTailMessages.length > 0) { return { - type: 'replace', + type: 'commit', messages: currentMessages, - durableMessages: durableReplacement.messages, - metadata: durableReplacement.metadata, - ...(durableReplacement.afterCommit ? { afterCommit: durableReplacement.afterCommit } : {}), + tailMessages: pendingTailMessages, + ...(durableReplacement ? { durableReplacement } : {}), }; } return requestMessagesReplaced @@ -410,6 +439,30 @@ async function runBeforeLLM( : { type: 'unchanged', messages }; } +interface DurableReplacement { + messages: AgentMessage[]; + metadata: PiBeforeLlmCallReplaceMetadata; + afterCommit?: PiAfterLlmReplacementCommit; +} + +function logAppendRegistered( + turn: turnState, + phase: 'initial' | 'iteration', + reason: string, + placement: 'tail' | 'before-current-user', +): void { + turn.logger.info( + { + session_id: turn.input.sessionId, + turn_id: turn.input.turnId, + phase, + reason, + placement, + }, + '[pi-turn-runner] beforeLlmCall registered hidden context message', + ); +} + function captureInitialCurrentUser( messages: readonly AgentMessage[], phase: 'initial' | 'iteration', @@ -452,6 +505,16 @@ function placeBeforeCurrentUser(input: { return 'before-current-user could not locate one current real user'; } const sourceIndex = input.initialMessages.length - 1; + // `currentUserSourceIndex` is reserved for placement (validReplacement + // rejects it from hooks), so its presence means an earlier hook in this + // pipeline already placed a marker and moved the exact current user to the + // tail of both views. + if (input.durableReplacement.metadata.currentUserSourceIndex !== undefined) { + return placeAfterEarlierPlacement( + { ...input, durableReplacement: input.durableReplacement }, + currentUser, + ); + } const sourceIndexes = input.durableReplacement.metadata.replacementSourceIndexes; const durableMatches = input.durableReplacement.messages.flatMap((message, index) => message === currentUser || sourceIndexes?.[index] === sourceIndex ? [index] : [], @@ -487,6 +550,65 @@ function placeBeforeCurrentUser(input: { }; } +/** + * Places another marker once an earlier placement already put the exact + * current user at the durable tail. Markers keep registration order, all + * before the current user. In the request view the user is matched by + * identity; if a request-only hook rewrote the user (e.g. a media guard + * replacing an image block), the first placement's invariant — the current + * user is the last request message — locates it instead. + */ +function placeAfterEarlierPlacement( + input: { + readonly currentMessages: readonly AgentMessage[]; + readonly durableReplacement: { + readonly messages: AgentMessage[]; + readonly metadata: PiBeforeLlmCallReplaceMetadata; + }; + readonly marker: PiBeforeLlmCallAppendMessage; + }, + currentUser: AgentMessage, +): + | string + | { + readonly type: 'placed'; + readonly messages: AgentMessage[]; + readonly durableReplacement: { + readonly messages: AgentMessage[]; + readonly metadata: PiBeforeLlmCallReplaceMetadata; + }; + } { + const durable = input.durableReplacement.messages; + if (durable.at(-1) !== currentUser || durable.indexOf(currentUser) !== durable.length - 1) { + return 'before-current-user lost the placed current user'; + } + const identityMatches = input.currentMessages.flatMap((message, index) => + message === currentUser ? [index] : [], + ); + if (identityMatches.length > 1) return 'before-current-user located multiple request users'; + const lastIndex = input.currentMessages.length - 1; + const requestUserIndex = + identityMatches[0] ?? (input.currentMessages[lastIndex]?.role === 'user' ? lastIndex : -1); + if (requestUserIndex < 0) return 'before-current-user lost the placed request user'; + const durableMessages = [...durable.slice(0, -1), input.marker, currentUser]; + return { + type: 'placed', + messages: [ + ...input.currentMessages.slice(0, requestUserIndex), + input.marker, + ...input.currentMessages.slice(requestUserIndex), + ], + durableReplacement: { + messages: durableMessages, + metadata: { + ...input.durableReplacement.metadata, + messagesAfter: durableMessages.length, + placedMarkerCount: (input.durableReplacement.metadata.placedMarkerCount ?? 1) + 1, + }, + }, + }; +} + function withCurrentUserTailMetadata( metadata: PiBeforeLlmCallReplaceMetadata, removedIndex: number | undefined, @@ -501,6 +623,7 @@ function withCurrentUserTailMetadata( ...metadata, messagesAfter, currentUserSourceIndex, + placedMarkerCount: 1, ...(replacementSourceIndexes === undefined ? {} : { replacementSourceIndexes }), }; } @@ -526,6 +649,9 @@ function validReplacement(metadata: PiBeforeLlmCallReplaceMetadata): string | un if (Object.hasOwn(metadata, 'currentUserSourceIndex')) { return 'currentUserSourceIndex is reserved for placement'; } + if (Object.hasOwn(metadata, 'placedMarkerCount')) { + return 'placedMarkerCount is reserved for placement'; + } if (!metadata.replacementId) return 'replacementId is required'; if (!metadata.strategyVersion) return 'strategyVersion is required'; if (typeof metadata.summary !== 'string') return 'summary is required'; diff --git a/packages/local-runtime-v2/src/service/session-system/messages/history/canonical-history-provider.ts b/packages/local-runtime-v2/src/service/session-system/messages/history/canonical-history-provider.ts index ee1543eca..96760bcc3 100644 --- a/packages/local-runtime-v2/src/service/session-system/messages/history/canonical-history-provider.ts +++ b/packages/local-runtime-v2/src/service/session-system/messages/history/canonical-history-provider.ts @@ -98,6 +98,8 @@ interface SessionCanonicalHistoryCompactionChangeBase extends Omit< readonly summary: string; /** Source identity for a structurally preserved real-user tail. */ readonly currentUserSourceIndex?: number; + /** Hidden markers placed directly before that tail; absent means 1. */ + readonly placedMarkerCount?: number; } export type SessionCanonicalHistoryCompactionChange = SessionCanonicalHistoryCompactionChangeBase & @@ -747,24 +749,32 @@ function readCurrentUserTail( if (!Number.isSafeInteger(sourceIndex) || sourceIndex < 0 || sourceIndex >= active.length) { throw new TypeError('Canonical current-user tail source lineage is invalid.'); } - if (change.messages.length < 2) { + const markerCount = change.placedMarkerCount ?? 1; + if (!Number.isSafeInteger(markerCount) || markerCount < 1) { + throw new TypeError('Canonical current-user tail marker count is invalid.'); + } + const tailLength = markerCount + 1; + if (change.messages.length < tailLength) { throw new TypeError('Canonical current-user tail is incomplete.'); } - const marker = change.messages.at(-2); + const baseLength = change.messages.length - tailLength; + const markers = change.messages.slice(baseLength, -1); const userMessage = change.messages.at(-1); const source = requireExternalCurrentUserSource(active[sourceIndex]); - const normalizedMarker = requireHiddenTypedCustom(marker); + const normalizedMarkers = markers.map(requireHiddenTypedCustom); const normalizedUser = requireExactCurrentUserMessage(userMessage, source); return { - baseMessages: change.messages.slice(0, -2), + baseMessages: change.messages.slice(0, baseLength), user: source, envelopes: [ - envelope({ - change, - message: normalizedMarker, - turnId: `${change.turnId}:compaction`, - index: change.messages.length - 2, - }), + ...normalizedMarkers.map((message, offset) => + envelope({ + change, + message, + turnId: `${change.turnId}:compaction`, + index: baseLength + offset, + }), + ), { ...withoutHistoryArtifact(source), message: normalizedUser }, ], }; diff --git a/packages/local-runtime-v2/src/service/turn-system/agent-host/compaction/context-compaction.test.ts b/packages/local-runtime-v2/src/service/turn-system/agent-host/compaction/context-compaction.test.ts index 42b58c610..f560a614c 100644 --- a/packages/local-runtime-v2/src/service/turn-system/agent-host/compaction/context-compaction.test.ts +++ b/packages/local-runtime-v2/src/service/turn-system/agent-host/compaction/context-compaction.test.ts @@ -1448,30 +1448,28 @@ describe('ContextCompactor result identity', () => { }); }); -describe('bounded compaction replacement metadata', () => { - const rawHistory = [{ role: 'user' as const, content: 'SENTINEL_RAW_H0', timestamp: 1 }]; - const replacement = [ - { role: 'user' as const, content: 'SENTINEL_REPLACEMENT_BODY', timestamp: 2 }, - ]; +const rawHistory = [{ role: 'user' as const, content: 'SENTINEL_RAW_H0', timestamp: 1 }]; +const replacement = [{ role: 'user' as const, content: 'SENTINEL_REPLACEMENT_BODY', timestamp: 2 }]; - function completed(method: 'tool_archive' | 'tool_trim' | 'llm_checkpoint') { - return { - status: 'completed' as const, - compactionId: `${method}-1`, - strategyVersion: 'context-compaction-v2', - method, - ...(method === 'llm_checkpoint' ? { summary: 'checkpoint正文' } : {}), - ...(method !== 'llm_checkpoint' ? { replacementSourceIndexes: [0] } : {}), - replacementMessages: replacement, - messagesBefore: 1, - messagesAfter: 1, - tokensBefore: 1_000, - tokensAfter: 400, - serializedBytesBefore: 4_000, - serializedBytesAfter: 1_600, - }; - } +function completed(method: 'tool_archive' | 'tool_trim' | 'llm_checkpoint') { + return { + status: 'completed' as const, + compactionId: `${method}-1`, + strategyVersion: 'context-compaction-v2', + method, + ...(method === 'llm_checkpoint' ? { summary: 'checkpoint正文' } : {}), + ...(method !== 'llm_checkpoint' ? { replacementSourceIndexes: [0] } : {}), + replacementMessages: replacement, + messagesBefore: 1, + messagesAfter: 1, + tokensBefore: 1_000, + tokensAfter: 400, + serializedBytesBefore: 4_000, + serializedBytesAfter: 1_600, + }; +} +describe('bounded compaction replacement metadata', () => { it.each(['tool_archive', 'tool_trim', 'llm_checkpoint'] as const)( 'maps automatic %s through Pi-compatible but body-free metadata', async (method) => { @@ -1597,7 +1595,9 @@ describe('bounded compaction replacement metadata', () => { ).toBeUndefined(); }, ); +}); +describe('bounded compaction current-user tail metadata', () => { it.each(['llm_checkpoint', 'tool_archive', 'tool_trim'] as const)( 'reads an exact current-user tail source for %s', (method) => { @@ -1620,6 +1620,83 @@ describe('bounded compaction replacement metadata', () => { }, ); + it.each(['tool_archive', 'tool_trim'] as const)( + 'accepts a %s commit with two placed markers before the current user', + (method) => { + const metadata = { + ...createCompactionReplaceMetadata(completed(method) as never, 'initial', 'attempt-2m'), + messagesBefore: 2, + messagesAfter: 4, + currentUserSourceIndex: 1, + placedMarkerCount: 2, + replacementSourceIndexes: [0], + }; + + expect(readCompactionCommitMetadata(metadata)).toMatchObject({ + method, + replacementSourceIndexes: [0], + currentUserSourceIndex: 1, + placedMarkerCount: 2, + }); + // Legacy single-marker reading (count absent) does not fit two markers. + const { placedMarkerCount: _omitted, ...legacy } = metadata; + expect(readCompactionCommitMetadata(legacy)).toBeUndefined(); + expect(readCompactionCommitMetadata({ ...metadata, placedMarkerCount: 1 })).toBeUndefined(); + expect(readCompactionCommitMetadata({ ...metadata, placedMarkerCount: 3 })).toBeUndefined(); + }, + ); + + it('accepts an llm_checkpoint commit with two placed markers before the current user', () => { + const metadata = { + ...createCompactionReplaceMetadata( + completed('llm_checkpoint') as never, + 'initial', + 'attempt-2m', + ), + messagesBefore: 2, + messagesAfter: 4, + currentUserSourceIndex: 1, + placedMarkerCount: 2, + }; + + expect(readCompactionCommitMetadata(metadata)).toMatchObject({ + method: 'llm_checkpoint', + currentUserSourceIndex: 1, + placedMarkerCount: 2, + }); + expect(readCompactionCommitMetadata({ ...metadata, messagesAfter: 6 })).toBeUndefined(); + expect(readCompactionCommitMetadata({ ...metadata, placedMarkerCount: 4 })).toBeUndefined(); + }); + + it.each([ + ['zero', 0], + ['negative', -1], + ['fractional', 1.5], + ['non-numeric', '2'], + ['unsafe', Number.MAX_SAFE_INTEGER + 1], + ])('rejects a %s placedMarkerCount', (_label, placedMarkerCount) => { + const metadata = { + ...createCompactionReplaceMetadata( + completed('llm_checkpoint') as never, + 'initial', + 'attempt-bad', + ), + messagesBefore: 2, + messagesAfter: 3, + currentUserSourceIndex: 1, + placedMarkerCount, + }; + expect(readCompactionCommitMetadata(metadata)).toBeUndefined(); + }); + + it('rejects placedMarkerCount without a current-user tail source', () => { + const metadata = { + ...createCompactionReplaceMetadata(completed('tool_trim') as never, 'initial', 'attempt-x'), + placedMarkerCount: 1, + }; + expect(readCompactionCommitMetadata(metadata)).toBeUndefined(); + }); + it.each([ ['missing method', { compactionMethod: undefined }], ['unknown method', { compactionMethod: 'recursive_reduce' }], diff --git a/packages/local-runtime-v2/src/service/turn-system/agent-host/compaction/context-compaction.ts b/packages/local-runtime-v2/src/service/turn-system/agent-host/compaction/context-compaction.ts index 6ab3e3db7..86550d32f 100644 --- a/packages/local-runtime-v2/src/service/turn-system/agent-host/compaction/context-compaction.ts +++ b/packages/local-runtime-v2/src/service/turn-system/agent-host/compaction/context-compaction.ts @@ -892,37 +892,46 @@ type CompactionCommitMetadata = readonly summary: string; readonly replacementSourceIndexes: readonly number[]; readonly currentUserSourceIndex?: number; + readonly placedMarkerCount?: number; } | { readonly compactionId: string; readonly method: 'llm_checkpoint'; readonly summary: string; readonly currentUserSourceIndex?: number; + readonly placedMarkerCount?: number; }; +/** Placement-owned current-user tail: source lineage plus the hidden markers before it. */ +interface CurrentUserTailMetadata { + readonly currentUserSourceIndex: number; + readonly placedMarkerCount: number; +} + export function readCompactionCommitMetadata(value: unknown): CompactionCommitMetadata | undefined { const lifecycle = readCompactionLifecycleMetadata(value); if (!lifecycle || !isRecord(value)) return undefined; const summary = Reflect.get(value, 'summary'); if (typeof summary !== 'string') return undefined; - const currentUserSourceIndex = readOptionalCurrentUserSourceIndex(value, lifecycle); - if (currentUserSourceIndex === null) return undefined; + const tail = readOptionalCurrentUserTail(value, lifecycle); + if (tail === null) return undefined; return lifecycle.method === 'llm_checkpoint' - ? readCheckpointCommitMetadata(lifecycle, summary, currentUserSourceIndex) - : readToolResultCommitMetadata(value, lifecycle, summary, currentUserSourceIndex); + ? readCheckpointCommitMetadata(lifecycle, summary, tail) + : readToolResultCommitMetadata(value, lifecycle, summary, tail); } function readCheckpointCommitMetadata( lifecycle: ContextCompactionLifecycleMetadata, summary: string, - currentUserSourceIndex: number | undefined, + tail: CurrentUserTailMetadata | undefined, ): CompactionCommitMetadata | undefined { - if (currentUserSourceIndex !== undefined && lifecycle.messagesAfter !== 3) return undefined; + // Native summary, then the exact markers + external user tail. + if (tail !== undefined && lifecycle.messagesAfter !== 2 + tail.placedMarkerCount) return undefined; return Object.freeze({ compactionId: lifecycle.compactionId, method: 'llm_checkpoint', summary, - ...(currentUserSourceIndex === undefined ? {} : { currentUserSourceIndex }), + ...currentUserTailFields(tail), }); } @@ -930,12 +939,14 @@ function readToolResultCommitMetadata( value: Readonly>, lifecycle: ContextCompactionLifecycleMetadata, summary: string, - currentUserSourceIndex: number | undefined, + tail: CurrentUserTailMetadata | undefined, ): CompactionCommitMetadata | undefined { if (lifecycle.method === 'llm_checkpoint') return undefined; const replacementSourceIndexes = Reflect.get(value, 'replacementSourceIndexes'); const expectedSourceCount = - currentUserSourceIndex === undefined ? lifecycle.messagesAfter : lifecycle.messagesAfter - 2; + tail === undefined + ? lifecycle.messagesAfter + : lifecycle.messagesAfter - 1 - tail.placedMarkerCount; if ( !Array.isArray(replacementSourceIndexes) || expectedSourceCount < 0 || @@ -945,7 +956,7 @@ function readToolResultCommitMetadata( Number.isSafeInteger(index) && index >= 0 && index < lifecycle.messagesBefore && - index !== currentUserSourceIndex && + index !== tail?.currentUserSourceIndex && (position === 0 || index > replacementSourceIndexes[position - 1]), ) ) { @@ -956,24 +967,51 @@ function readToolResultCommitMetadata( method: lifecycle.method, summary, replacementSourceIndexes: Object.freeze([...replacementSourceIndexes]) as readonly number[], - ...(currentUserSourceIndex === undefined ? {} : { currentUserSourceIndex }), + ...currentUserTailFields(tail), }); } -function readOptionalCurrentUserSourceIndex( +function currentUserTailFields( + tail: CurrentUserTailMetadata | undefined, +): Partial { + return tail === undefined + ? {} + : { + currentUserSourceIndex: tail.currentUserSourceIndex, + placedMarkerCount: tail.placedMarkerCount, + }; +} + +/** + * Reads the placement-owned tail lineage. `placedMarkerCount` defaults to 1 so + * metadata persisted before multi-marker placement stays readable; it is + * invalid without `currentUserSourceIndex`. Returns null for invalid input. + */ +function readOptionalCurrentUserTail( value: Readonly>, lifecycle: { readonly messagesBefore: number }, -): number | undefined | null { - if (!Object.hasOwn(value, 'currentUserSourceIndex')) return undefined; +): CurrentUserTailMetadata | undefined | null { + const hasMarkerCount = Object.hasOwn(value, 'placedMarkerCount'); + if (!Object.hasOwn(value, 'currentUserSourceIndex')) return hasMarkerCount ? null : undefined; const index = Reflect.get(value, 'currentUserSourceIndex'); - return typeof index === 'number' && - Number.isSafeInteger(index) && - index >= 0 && - index < lifecycle.messagesBefore - ? index + if (!isSourceIndexInRange(index, lifecycle.messagesBefore)) return null; + if (!hasMarkerCount) return { currentUserSourceIndex: index, placedMarkerCount: 1 }; + const count = Reflect.get(value, 'placedMarkerCount'); + return isPositiveSafeInteger(count) + ? { currentUserSourceIndex: index, placedMarkerCount: count } : null; } +function isSourceIndexInRange(value: unknown, messagesBefore: number): value is number { + return ( + typeof value === 'number' && Number.isSafeInteger(value) && value >= 0 && value < messagesBefore + ); +} + +function isPositiveSafeInteger(value: unknown): value is number { + return typeof value === 'number' && Number.isSafeInteger(value) && value > 0; +} + export function readCompactionAttemptId(value: unknown): string | undefined { return isRecord(value) ? readNonEmpty(value, 'compactionAttemptId') : undefined; } diff --git a/packages/local-runtime-v2/src/service/turn-system/agent-host/execution/executor.test.ts b/packages/local-runtime-v2/src/service/turn-system/agent-host/execution/executor.test.ts index ec3a89c4c..5d6829317 100644 --- a/packages/local-runtime-v2/src/service/turn-system/agent-host/execution/executor.test.ts +++ b/packages/local-runtime-v2/src/service/turn-system/agent-host/execution/executor.test.ts @@ -64,7 +64,6 @@ import type { AgentHostFileChangeObservation, LocalTurnExecutionInput, } from "../runner/contracts.js"; -import { createExecutionBudgetReminder } from "../runner/execution-budget-reminder.js"; import type { AgentEventDelivery } from "../events/contracts.js"; import type { CanonicalHistoryStore } from "../history/contracts.js"; import { copyCanonicalHistoryForPiCompatibility } from "../history/canonical-history-validation.js"; @@ -4377,6 +4376,8 @@ describe("LocalRuntimeTurnExecutor session id reminder context", () => { ); }); +const EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE = "execution_budget_reminder"; + describe("LocalRuntimeTurnExecutor execution budget reminders", () => { it.each([ { @@ -4404,56 +4405,52 @@ describe("LocalRuntimeTurnExecutor execution budget reminders", () => { ], }, ] as const)( - "reminds a budgeted run ($durationMs ms) only at thresholds without changing history or the system prompt", + "reminds a budgeted run ($durationMs ms) only at thresholds without changing the system prompt", async ({ durationMs, steps }) => { let nowMs = 1_000; const executorOptions = options(async (runInput) => { - const messages = [ - { role: "user" as const, content: "hello", timestamp: 1 }, - ]; // The budget reminder is registered immediately before the background // cadence reminder, so it is the second-to-last host hook. const hook = runInput.hooks?.beforeLlmCallHook?.at(-2); - const hookInput = { + // PiTurnRunner commits every appended marker before the next request. + let messages: PiBeforeLlmCallHookInput["messages"] = [ + { role: "user", content: "hello", timestamp: 1 }, + ]; + const hookInput = (phase: "initial" | "iteration") => ({ sessionId: runInput.sessionId, turnId: runInput.turnId, - phase: "initial" as const, + phase, messages, canonicalMessages: messages, model: runInput.llm.model, thinkingLevel: "off" as const, - }; + }); for (const [elapsedMs, expected] of steps) { nowMs = 1_000 + elapsedMs; - const decision = await hook?.({ - ...hookInput, - phase: elapsedMs === 0 ? "initial" : "iteration", - }); + const decision = await hook?.( + hookInput(elapsedMs === 0 ? "initial" : "iteration"), + ); if (expected === undefined) { expect(decision).toBeUndefined(); continue; } expect(decision).toMatchObject({ - type: "replaceRequestMessages", - reason: "execution-budget-context", - messages: [ - messages[0], - { - role: "user", - content: expect.stringContaining( - `Execution time remaining: ${expected}`, - ), - }, - ], + type: "appendMessage", + reason: EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE, + message: { + role: "custom", + customType: EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE, + display: false, + content: expect.stringContaining( + `Execution time remaining: ${expected}`, + ), + }, }); + if (decision?.type === "appendMessage") + messages = [...messages, decision.message]; } nowMs = 1_000 + durationMs; - expect( - await hook?.({ ...hookInput, phase: "iteration" }), - ).toBeUndefined(); - expect(messages).toEqual([ - { role: "user", content: "hello", timestamp: 1 }, - ]); + expect(await hook?.(hookInput("iteration"))).toBeUndefined(); expect(runInput.systemPrompt).toBe("extension system\n\nbase system"); await runInput.eventWriter.pushRuntime( terminalEvent(RuntimeEventStatus.COMPLETED), @@ -4472,66 +4469,9 @@ describe("LocalRuntimeTurnExecutor execution budget reminders", () => { executionDeadlineAtMs: 1_000 + durationMs, }, }); - - expect(input.onHistoryChanged).not.toHaveBeenCalled(); }, ); - it("retries a threshold reminder that admission rejected", async () => { - let nowMs = 1_000; - let admit = false; - const hook = createExecutionBudgetReminder( - 1_000 + 1_800_000, - () => nowMs, - () => admit, - )!; - const messages = [{ role: "user" as const, content: "hello", timestamp: 1 }]; - const hookInput = { - sessionId: "session-1", - turnId: "turn-1", - phase: "initial" as const, - messages, - canonicalMessages: messages, - model: {} as never, - thinkingLevel: "off" as const, - }; - expect(await hook(hookInput)).toBeUndefined(); - admit = true; - nowMs = 2_000; - expect(await hook(hookInput)).toMatchObject({ - messages: [ - messages[0], - { - content: - "Execution time remaining: about 29 minutes of 30 minutes total. Model generation, tools, and waiting all count against it.", - }, - ], - }); - nowMs = 1_000 + 1_800_000 - 899_000; - expect(await hook(hookInput)).toMatchObject({ - messages: [ - messages[0], - { - content: - "Execution time remaining: about 14 minutes of 30 minutes total.", - }, - ], - }); - // Jumping straight past 25% and 10% emits once for the deepest threshold. - nowMs = 1_000 + 1_800_000 - 179_000; - expect(await hook(hookInput)).toMatchObject({ - messages: [ - messages[0], - { - content: - "Execution time remaining: about 2 minutes of 30 minutes total.", - }, - ], - }); - nowMs += 1_000; - expect(await hook(hookInput)).toBeUndefined(); - }); - it("samples budget context while preserving existing request observers", async () => { let nowMs = 1_000; const existingObserver = vi.fn(); @@ -4550,12 +4490,19 @@ describe("LocalRuntimeTurnExecutor execution budget reminders", () => { thinkingLevel: "off" as const, }; const initial = await runInput.hooks!.beforeLlmCallHook!.at(-2)!(input); - if (initial?.type !== "replaceRequestMessages") + if (initial?.type !== "appendMessage") throw new Error("Expected initial reminder"); + // PiTurnRunner commits the marker before the request and later samples. + const committed = [...messages, initial.message]; + const next = { + ...input, + messages: committed, + canonicalMessages: committed, + }; for (const observer of runInput.hooks!.onLlmCallPreparedHook!) { await observer({ ...input, - messages: initial.messages, + messages: committed, scope: "agent", systemPrompt: "", tools: [], @@ -4569,15 +4516,17 @@ describe("LocalRuntimeTurnExecutor execution budget reminders", () => { // 99 seconds later no threshold is crossed, so no tokens are spent. nowMs = 100_000; expect( - await runInput.hooks!.beforeLlmCallHook!.at(-2)!(input), + await runInput.hooks!.beforeLlmCallHook!.at(-2)!(next), ).toBeUndefined(); expect(logger.info).toHaveBeenCalledOnce(); nowMs = 301_000; - const next = await runInput.hooks!.beforeLlmCallHook!.at(-2)!(input); - if (next?.type !== "replaceRequestMessages") + const threshold = await runInput.hooks!.beforeLlmCallHook!.at(-2)!(next); + if (threshold?.type !== "appendMessage") throw new Error("Expected time context"); - expect(next.messages.at(-1)).toMatchObject({ - content: expect.stringContaining("about 5 minutes of 10 minutes total."), + expect(threshold.message).toMatchObject({ + content: expect.stringContaining( + "about 5 minutes of 10 minutes total.", + ), }); await runInput.eventWriter.pushRuntime( terminalEvent(RuntimeEventStatus.COMPLETED), @@ -4600,9 +4549,10 @@ describe("LocalRuntimeTurnExecutor execution budget reminders", () => { }, request: { ...input.request, executionDeadlineAtMs: 601_000 }, }); - expect(input.onHistoryChanged).not.toHaveBeenCalled(); }); +}); +describe("LocalRuntimeTurnExecutor execution budget reminder registration", () => { it("does not add budget context to unbudgeted or already aborted runs", async () => { for (const aborted of [false, true]) { const controller = new AbortController(); @@ -4641,12 +4591,10 @@ describe("LocalRuntimeTurnExecutor execution budget reminders", () => { } }); - it("registers the request-only budget reminder before the appending background reminder", async () => { - // runBeforeLLM ends the before-LLM decision pipeline as soon as a hook - // returns appendMessage. The background cadence reminder appends, so a - // budget reminder registered after it would be silently skipped for every - // request where the background reminder fires. The relative order is the - // only thing keeping the request-only reminder observable. + it("registers the budget reminder before the background reminder", async () => { + // runBeforeLLM accumulates appendMessage decisions, so this order only + // fixes where the durable budget marker sits relative to the durable + // background marker, identically in the request and in history. const backgroundHook = vi.fn(() => undefined); const executorOptions = options(async (runInput) => { const messages = [ @@ -4668,8 +4616,8 @@ describe("LocalRuntimeTurnExecutor execution budget reminders", () => { ); const budgetIndex = decisions.findIndex( (decision) => - decision?.type === "replaceRequestMessages" && - decision.reason === "execution-budget-context", + decision?.type === "appendMessage" && + decision.reason === EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE, ); const backgroundIndex = hooks.indexOf(backgroundHook); expect(budgetIndex).toBeGreaterThanOrEqual(0); @@ -4704,7 +4652,7 @@ describe("LocalRuntimeTurnExecutor budget with durable reminders", () => { { admission: "allow", mode: "delayed" }, { admission: "allow", mode: "observed" }, ] as const)( - "preserves durable reminders with request-only time context ($mode, $admission)", + "keeps background reminders alongside the durable budget reminder ($mode, $admission)", async ({ admission, mode }) => { let nowMs = 1_000; const marker = { @@ -4743,11 +4691,28 @@ describe("LocalRuntimeTurnExecutor budget with durable reminders", () => { }; }, }); - const laterHook = vi.fn(replacement); + // runBeforeLLM accumulates appends, so this later hook runs in every + // mode and sees the registered markers. It replaces durable history only + // when nothing is pending, because a durable replacement after a + // registered append fails closed (production never orders them so). + const laterSawMarker: boolean[] = []; + const laterHook = vi.fn((hookInput: PiBeforeLlmCallHookInput) => { + const sawBackground = hookInput.messages.some( + (message) => message === marker, + ); + laterSawMarker.push(sawBackground); + const pending = + sawBackground || + hookInput.messages.some( + (message) => + message.role === "custom" && + message.customType === EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE, + ); + return pending ? { type: "continue" as const } : replacement(hookInput); + }); const canAppend = vi.fn((hookInput: PiBeforeLlmCallHookInput) => { - // The request-only reminder runs before the appending background - // reminder and before any post-commit callback, so it samples the - // pre-append request view. + // The budget reminder runs before the background reminder and before + // any post-commit callback, so it samples the pre-append request view. expect(hookInput.messages.some((message) => message === marker)).toBe( false, ); @@ -4821,28 +4786,29 @@ describe("LocalRuntimeTurnExecutor budget with durable reminders", () => { expect(providerContexts).toEqual(preparedContexts); const requestContext = providerContexts[0] ?? ""; expect(requestContext.includes(marker.content)).toBe(mode !== "defer"); - // A later durable replaceMessages hook rebuilds the request from canonical - // history, which drops an earlier request-only replacement ('defer' runs - // laterHook because the background marker cannot be placed). - const keepsRequestContext = admission === "allow" && mode !== "defer"; + // The budget marker is durable, so no later hook can drop it again. + const budgeted = admission === "allow"; expect( requestContext.match(/Execution time remaining/g) ?? [], - ).toHaveLength(keepsRequestContext ? 1 : 0); - if (keepsRequestContext) { + ).toHaveLength(budgeted ? 1 : 0); + if (budgeted) { expect(requestContext).toContain( "Execution time remaining: about 120 seconds of 120 seconds total.", ); } expect(canAppend).toHaveBeenCalledOnce(); expect(backgroundHook).toHaveBeenCalledOnce(); - expect(laterHook).toHaveBeenCalledTimes(mode === "defer" ? 1 : 0); + expect(laterHook).toHaveBeenCalledOnce(); + expect(laterSawMarker).toEqual([mode !== "defer"]); const historyChanges = vi.mocked(input.onHistoryChanged).mock.calls; expect(historyChanges.length).toBeGreaterThan(0); const history = historyChanges.flatMap(([change]) => change.messages); expect(JSON.stringify(history).includes(marker.content)).toBe( mode !== "defer", ); - expect(JSON.stringify(history)).not.toContain("Execution time remaining"); + expect( + JSON.stringify(history).match(/Execution time remaining/g) ?? [], + ).toHaveLength(budgeted ? 1 : 0); }, ); }); diff --git a/packages/local-runtime-v2/src/service/turn-system/agent-host/execution/executor.ts b/packages/local-runtime-v2/src/service/turn-system/agent-host/execution/executor.ts index 7911e6339..59eb34ea5 100644 --- a/packages/local-runtime-v2/src/service/turn-system/agent-host/execution/executor.ts +++ b/packages/local-runtime-v2/src/service/turn-system/agent-host/execution/executor.ts @@ -777,11 +777,11 @@ function pluginHookTranscriptCleanup( } /** - * Order is load-bearing: `runBeforeLLM` ends the before-LLM decision pipeline - * as soon as a hook returns `appendMessage`, so a request-only reminder placed - * after an appending hook is silently skipped for that request. The budget - * reminder is request-only and therefore always precedes the background - * cadence reminder, which appends. + * Both reminders append durable hidden markers and `runBeforeLLM` accumulates + * them, so neither can starve the other. The order fixes where the budget + * marker sits relative to the background marker, identically in the request + * and in history, and lets the budget admission measure the request before + * the background marker is registered. */ function afterCompactionReminderHooks( budgetReminder: PiBeforeLlmCallHook | undefined, diff --git a/packages/local-runtime-v2/src/service/turn-system/agent-host/runner/execution-budget-reminder.test.ts b/packages/local-runtime-v2/src/service/turn-system/agent-host/runner/execution-budget-reminder.test.ts new file mode 100644 index 000000000..6158223dc --- /dev/null +++ b/packages/local-runtime-v2/src/service/turn-system/agent-host/runner/execution-budget-reminder.test.ts @@ -0,0 +1,403 @@ +import type { AgentMessage } from '@earendil-works/pi-agent-core'; +import type { + PiBeforeLlmCallAppendMessage, + PiBeforeLlmCallHook, + PiBeforeLlmCallHookDecision, + PiBeforeLlmCallHookInput, +} from '@mavis/agent-core/pi-turn-runner'; +import { resolveCompactionTokenBudget } from '@mavis/context-manager'; +import { describe, expect, it, vi } from 'vitest'; + +import { createAutomaticFootprintMeasurer } from '../../compaction/automatic-context-compactor.js'; +import { ContextUsageAnchorState } from '../../compaction/execution/usage-anchor.js'; +import { fitsReminderInFinalRequest } from '../../execution/reminder/reminder-admission.js'; +import { validateCanonicalHistoryMessages } from '../history/canonical-history-validation.js'; +import { createExecutionBudgetReminder } from './execution-budget-reminder.js'; + +const EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE = 'execution_budget_reminder'; + +/** The reminder is durable hidden context: it must only ask for a tail append. */ +function appended( + decision: PiBeforeLlmCallHookDecision | undefined, +): PiBeforeLlmCallAppendMessage | undefined { + if (!decision) return undefined; + if (decision.type !== 'appendMessage') { + throw new Error(`Expected appendMessage, received ${decision.type}`); + } + expect(decision.reason).toBe(EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE); + expect(decision).not.toHaveProperty('placement'); + return decision.message; +} + +/** Runs one request and commits its marker the way PiTurnRunner does. */ +async function request( + reminder: PiBeforeLlmCallHook, + input: PiBeforeLlmCallHookInput, +): Promise { + const marker = appended(await reminder(input)); + if (marker) { + input.messages = [...input.messages, marker]; + input.canonicalMessages = [...input.canonicalMessages, marker]; + } + return marker; +} + +function reminderText(text: string): string { + return `Execution time remaining: ${text}`; +} + +function hookInput(): PiBeforeLlmCallHookInput { + const messages: AgentMessage[] = [{ role: 'user', content: 'hello', timestamp: 1 }]; + return { + sessionId: 'session-1', + turnId: 'turn-1', + phase: 'initial', + messages, + canonicalMessages: messages, + model: { + id: 'test-model', + name: 'test-model', + api: 'anthropic-messages', + provider: 'test', + baseUrl: 'https://example.invalid', + reasoning: false, + input: ['text'], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 200_000, + maxTokens: 8_192, + }, + thinkingLevel: 'off', + tools: [], + }; +} + +function tokenFullRequest(usageAnchor: ContextUsageAnchorState): PiBeforeLlmCallHookInput { + const input = hookInput(); + const { providerInputLimit } = resolveCompactionTokenBudget({ + contextWindow: input.model.contextWindow, + configuredMaxOutputTokens: input.model.maxTokens, + }); + const anchor: AgentMessage = { + role: 'assistant', + content: [{ type: 'text', text: 'usage anchor' }], + api: input.model.api, + provider: input.model.provider, + model: input.model.id, + responseId: 'budget-anchor', + timestamp: 2, + stopReason: 'stop', + usage: { + input: providerInputLimit - 10, + output: 10, + cacheRead: 0, + cacheWrite: 0, + totalTokens: providerInputLimit, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }; + const binding = usageAnchor.bind({ + scope: input.sessionId, + provider: input.model.provider, + api: input.model.api, + model: input.model.id, + tools: input.tools, + }); + expect( + usageAnchor.record({ + key: binding.key, + historyEpoch: binding.historyEpoch, + assistant: anchor, + }), + ).toBe(true); + input.messages = [anchor]; + input.canonicalMessages = input.messages; + expect( + createAutomaticFootprintMeasurer(input, usageAnchor).measure(input.messages).inputTokens, + ).toBe(providerInputLimit); + return input; +} + +function byteFullRequest(usageAnchor: ContextUsageAnchorState): PiBeforeLlmCallHookInput { + const input = hookInput(); + const footprint = createAutomaticFootprintMeasurer(input, usageAnchor).measure(input.messages); + input.maxSerializedInputBytes = footprint.serializedBytes; + return input; +} + +describe('Execution budget durable context', () => { + it('reminds only at budget thresholds and keeps every reminder in history', async () => { + let nowMs = 1_000; + const reminder = createExecutionBudgetReminder( + 601_000, + () => nowMs, + () => true, + )!; + const input = hookInput(); + const committed: PiBeforeLlmCallAppendMessage[] = []; + for (const [now, expected, level] of [ + [ + 1_000, + 'about 10 minutes of 10 minutes total. Model generation, tools, and waiting all count against it.', + 0, + ], + [100_000, undefined], + [301_000, 'about 5 minutes of 10 minutes total.', 1], + [400_000, undefined], + [451_000, 'about 2 minutes of 10 minutes total.', 2], + [481_000, 'about 120 seconds of 10 minutes total.', 4], + [490_000, 'about 110 seconds of 10 minutes total.', 4], + // Same 10-second step as the previous reminder: nothing new to say. + [491_000, undefined], + [495_000, 'about 100 seconds of 10 minutes total.', 4], + [600_999, 'under 10 seconds of 10 minutes total.', 4], + ] as const) { + nowMs = now; + const marker = await request(reminder, input); + if (expected === undefined) { + expect(marker).toBeUndefined(); + continue; + } + expect(marker).toEqual({ + role: 'custom', + customType: EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE, + content: reminderText(expected), + display: false, + details: { deadlineAtMs: 601_000, level }, + timestamp: now, + }); + committed.push(marker!); + } + expect(input.canonicalMessages).toEqual([hookInput().messages[0], ...committed]); + }); + + it('stores reminders that canonical history accepts', async () => { + const marker = appended( + await createExecutionBudgetReminder( + 601_000, + () => 1_000, + () => true, + )!(hookInput()), + ); + expect(() => validateCanonicalHistoryMessages([marker])).not.toThrow(); + }); + + it('announces the budget again once compaction summarized earlier reminders away', async () => { + let nowMs = 1_000; + const reminder = createExecutionBudgetReminder( + 601_000, + () => nowMs, + () => true, + )!; + const input = hookInput(); + expect(await request(reminder, input)).toBeDefined(); + nowMs = 301_000; + expect(await request(reminder, input)).toBeDefined(); + const summary: AgentMessage = { role: 'user', content: 'compacted summary', timestamp: 3 }; + input.messages = [summary]; + input.canonicalMessages = [summary]; + nowMs = 320_000; + expect((await request(reminder, input))?.content).toBe( + reminderText( + 'about 4 minutes of 10 minutes total. Model generation, tools, and waiting all count against it.', + ), + ); + nowMs = 330_000; + expect(await request(reminder, input)).toBeUndefined(); + }); + + it('ignores reminders of another budget and malformed reminder details', async () => { + const reminder = createExecutionBudgetReminder( + 601_000, + () => 2_000, + () => true, + )!; + const input = hookInput(); + const earlier = { + role: 'custom' as const, + customType: EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE, + content: reminderText('about 1 minutes of 5 minutes total.'), + display: false as const, + timestamp: 2, + }; + input.canonicalMessages = [ + ...input.canonicalMessages, + { ...earlier, details: { deadlineAtMs: 300_000, level: 4 } }, + { ...earlier, details: { deadlineAtMs: '601000', level: 4 } }, + earlier, + ]; + expect((await request(reminder, input))?.content).toContain( + 'Model generation, tools, and waiting all count against it.', + ); + }); +}); + +describe('Execution budget request admission', () => { + it.each([ + ['provider token limit', tokenFullRequest], + ['serialized byte limit', byteFullRequest], + ] as const)('preserves final request admission for its %s', async (_name, fullRequest) => { + const usageAnchor = new ContextUsageAnchorState(); + const reminder = createExecutionBudgetReminder( + 601_000, + () => 1_000, + (input, marker) => fitsReminderInFinalRequest(input, marker, usageAnchor), + )!; + expect(await reminder(fullRequest(usageAnchor))).toBeUndefined(); + expect(appended(await reminder(hookInput()))).toBeDefined(); + }); + + it.each(['deny', 'throw'] as const)( + 'retries a threshold reminder that admission skipped (%s)', + async (failure) => { + let nowMs = 1_000; + let skip = false; + const logger = { info: vi.fn() }; + const reminder = createExecutionBudgetReminder( + 601_000, + () => nowMs, + () => { + if (skip && failure === 'throw') throw new Error('Measurement unavailable'); + return !skip; + }, + logger, + )!; + const input = hookInput(); + expect(await request(reminder, input)).toBeDefined(); + skip = true; + nowMs = 301_000; + expect(await request(reminder, input)).toBeUndefined(); + expect(logger.info).toHaveBeenCalledOnce(); + skip = false; + nowMs = 311_000; + expect((await request(reminder, input))?.content).toBe( + reminderText('about 4 minutes of 10 minutes total.'), + ); + nowMs = 321_000; + expect(await request(reminder, input)).toBeUndefined(); + }, + ); + + it('keeps the first reminder pending until admitted and collapses skipped thresholds', async () => { + let nowMs = 1_000; + let admit = false; + const reminder = createExecutionBudgetReminder( + 1_000 + 1_800_000, + () => nowMs, + () => admit, + )!; + const input = hookInput(); + expect(await request(reminder, input)).toBeUndefined(); + admit = true; + nowMs = 2_000; + expect((await request(reminder, input))?.content).toBe( + reminderText( + 'about 29 minutes of 30 minutes total. Model generation, tools, and waiting all count against it.', + ), + ); + nowMs = 1_000 + 1_800_000 - 899_000; + expect((await request(reminder, input))?.content).toBe( + reminderText('about 14 minutes of 30 minutes total.'), + ); + // Jumping straight past 25% and 10% emits once for the deepest threshold. + nowMs = 1_000 + 1_800_000 - 179_000; + expect((await request(reminder, input))?.content).toBe( + reminderText('about 2 minutes of 30 minutes total.'), + ); + nowMs += 1_000; + expect(await request(reminder, input)).toBeUndefined(); + }); + + it('reports the sampled budget level in diagnostics', async () => { + let nowMs = 1_000; + const logger = { info: vi.fn() }; + const reminder = createExecutionBudgetReminder( + 601_000, + () => nowMs, + () => true, + logger, + )!; + const input = hookInput(); + expect(await request(reminder, input)).toMatchObject({ timestamp: 1_000 }); + expect(logger.info).toHaveBeenLastCalledWith( + { + session_id: 'session-1', + turn_id: 'turn-1', + sampled_at_ms: 1_000, + remaining_ms: 600_000, + total_ms: 600_000, + budget_level: 0, + }, + 'execution budget context sampled', + ); + nowMs = 451_000; + expect(await request(reminder, input)).toMatchObject({ timestamp: 451_000 }); + expect(logger.info).toHaveBeenLastCalledWith( + { + session_id: 'session-1', + turn_id: 'turn-1', + sampled_at_ms: 451_000, + remaining_ms: 150_000, + total_ms: 600_000, + budget_level: 2, + }, + 'execution budget context sampled', + ); + }); + + it('keeps admitted context when diagnostics fail', async () => { + const reminder = createExecutionBudgetReminder( + 601_000, + () => 1_000, + () => true, + { + info: () => { + throw new Error('Logger unavailable'); + }, + }, + )!; + expect(appended(await reminder(hookInput()))).toBeDefined(); + }); + + it('ignores a clock that moves backwards and isolates thresholds per turn', async () => { + let nowMs = 1_000; + const reminder = createExecutionBudgetReminder( + 601_000, + () => nowMs, + () => true, + )!; + const input = hookInput(); + expect(await request(reminder, input)).toBeDefined(); + nowMs = 500; + expect(await request(reminder, input)).toBeUndefined(); + const nextTurn = createExecutionBudgetReminder( + 602_000, + () => 2_000, + () => true, + )!; + expect((await request(nextTurn, input))?.content).toContain( + 'Model generation, tools, and waiting all count against it.', + ); + }); + + it('skips unbudgeted, aborted and expired requests', async () => { + expect( + createExecutionBudgetReminder( + undefined, + () => 1_000, + () => true, + ), + ).toBeUndefined(); + expect(createExecutionBudgetReminder(601_000)).toBeUndefined(); + let nowMs = 1_000; + const logger = { info: vi.fn() }; + const admission = vi.fn(() => true); + const reminder = createExecutionBudgetReminder(601_000, () => nowMs, admission, logger)!; + expect(await reminder({ ...hookInput(), signal: AbortSignal.abort() })).toBeUndefined(); + nowMs = 601_000; + expect(await reminder(hookInput())).toBeUndefined(); + nowMs += 1; + expect(await reminder(hookInput())).toBeUndefined(); + expect(admission).not.toHaveBeenCalled(); + expect(logger.info).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/local-runtime-v2/src/service/turn-system/agent-host/runner/execution-budget-reminder.ts b/packages/local-runtime-v2/src/service/turn-system/agent-host/runner/execution-budget-reminder.ts index f6e25a3a9..48dfe62ba 100644 --- a/packages/local-runtime-v2/src/service/turn-system/agent-host/runner/execution-budget-reminder.ts +++ b/packages/local-runtime-v2/src/service/turn-system/agent-host/runner/execution-budget-reminder.ts @@ -1,25 +1,47 @@ -import type { PiBeforeLlmCallHook, PiTurnRunnerLogger } from '@mavis/agent-core/pi-turn-runner'; +import type { + PiBeforeLlmCallAppendMessage, + PiBeforeLlmCallHook, + PiBeforeLlmCallHookInput, + PiTurnRunnerLogger, +} from '@mavis/agent-core/pi-turn-runner'; import type { ExecutionBudgetReminderAdmission } from '../execution/contracts.js'; +const EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE = 'execution_budget_reminder'; + +type HistoryMessage = PiBeforeLlmCallHookInput['canonicalMessages'][number]; + /** Remaining-budget fractions that each emit one reminder when first crossed. */ const BUDGET_FRACTION_THRESHOLDS = [0.5, 0.25, 0.1] as const; -/** Inside this window every request carries a (coarse) reminder. */ +/** Inside this window a reminder is emitted whenever its coarse text changes. */ const FINAL_WINDOW_MS = 120_000; const FINAL_WINDOW_LEVEL = BUDGET_FRACTION_THRESHOLDS.length + 1; +interface ExecutionBudgetReminderDetails { + /** Deadline the reminder was sampled for; separates the budgets of different Turns. */ + readonly deadlineAtMs: number; + readonly level: number; +} + +interface PreviousReminder { + readonly level: number; + readonly content: string; +} + /** - * Request-only context; the caller's existing cancellation timer remains - * authoritative. `replaceRequestMessages` keeps the marker out of durable - * history, so it must be registered before any hook that returns - * `appendMessage` — `runBeforeLLM` ends the pipeline on a successful append. + * Hidden durable context; the caller's existing cancellation timer remains + * authoritative. Each reminder is appended to canonical history, so every + * later provider request keeps it in place and the prompt-cache prefix stays + * append-only. * * To avoid spending tokens on every model request, the reminder is emitted - * only on the first request (announces the budget), when the remaining budget - * first drops to ≤50%, ≤25% and ≤10% of the total, and on every request inside - * the final two minutes. The total budget is measured from hook creation - * (turn preparation), since requests carry no start time. Precision is whole - * minutes outside the final window and 10-second steps inside it, so the text - * does not churn between otherwise identical requests. + * on the first request (announces the budget), when the remaining budget first + * drops to ≤50%, ≤25% and ≤10% of the total, and inside the final two minutes + * whenever its text changes. Progress is read back from the latest reminder of + * the same deadline in canonical history, so a repeated request does not + * duplicate it and a compaction that summarized it away announces the budget + * again. The total budget is measured from hook creation (turn preparation), + * since requests carry no start time. Precision is whole minutes outside the + * final window and 10-second steps inside it. */ export function createExecutionBudgetReminder( executionDeadlineAtMs: number | undefined, @@ -29,7 +51,6 @@ export function createExecutionBudgetReminder( ): PiBeforeLlmCallHook | undefined { if (executionDeadlineAtMs === undefined || !canAppend) return undefined; const totalMs = Math.max(0, executionDeadlineAtMs - nowMs()); - let emittedLevel: number | undefined; return (input) => { try { if (input.signal?.aborted) return undefined; @@ -37,15 +58,18 @@ export function createExecutionBudgetReminder( const remainingMs = executionDeadlineAtMs - now; if (remainingMs <= 0) return undefined; const level = budgetLevel(remainingMs, totalMs); - const first = emittedLevel === undefined; - if (!first && level !== FINAL_WINDOW_LEVEL && level <= emittedLevel!) return undefined; - const marker = { - role: 'user' as const, - content: formatBudgetReminder(remainingMs, totalMs, first), + const previous = latestReminder(input.canonicalMessages, executionDeadlineAtMs); + const content = formatBudgetReminder(remainingMs, totalMs, previous === undefined); + if (previous && !isReminderDue(previous, level, content)) return undefined; + const marker: PiBeforeLlmCallAppendMessage = { + role: 'custom', + customType: EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE, + content, + display: false, + details: { deadlineAtMs: executionDeadlineAtMs, level }, timestamp: now, }; if (!canAppend(input, marker)) return undefined; - emittedLevel = level; try { logger?.info?.( { @@ -62,9 +86,9 @@ export function createExecutionBudgetReminder( // Diagnostics must not affect the request context. } return { - type: 'replaceRequestMessages', - reason: 'execution-budget-context', - messages: [...input.messages, marker], + type: 'appendMessage', + reason: EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE, + message: marker, }; } catch { return undefined; @@ -72,6 +96,43 @@ export function createExecutionBudgetReminder( }; } +function isReminderDue(previous: PreviousReminder, level: number, content: string): boolean { + return level === FINAL_WINDOW_LEVEL ? content !== previous.content : level > previous.level; +} + +/** Latest reminder for this deadline; reminders of other budgets are ignored. */ +function latestReminder( + messages: readonly HistoryMessage[], + deadlineAtMs: number, +): PreviousReminder | undefined { + for (let index = messages.length - 1; index >= 0; index -= 1) { + const reminder = readReminder(messages[index]); + if (reminder?.details.deadlineAtMs === deadlineAtMs) { + return { level: reminder.details.level, content: reminder.content }; + } + } + return undefined; +} + +function readReminder( + message: HistoryMessage | undefined, +): { readonly content: string; readonly details: ExecutionBudgetReminderDetails } | undefined { + if ( + message?.role !== 'custom' || + message.customType !== EXECUTION_BUDGET_REMINDER_CUSTOM_TYPE || + typeof message.content !== 'string' + ) { + return undefined; + } + const details: unknown = message.details; + if (typeof details !== 'object' || details === null) return undefined; + const deadlineAtMs = Reflect.get(details, 'deadlineAtMs'); + const level = Reflect.get(details, 'level'); + return typeof deadlineAtMs === 'number' && typeof level === 'number' + ? { content: message.content, details: { deadlineAtMs, level } } + : undefined; +} + /** 0 = above 50%; 1..3 = crossed 50/25/10%; FINAL_WINDOW_LEVEL = final two minutes. */ function budgetLevel(remainingMs: number, totalMs: number): number { if (remainingMs <= FINAL_WINDOW_MS) return FINAL_WINDOW_LEVEL; @@ -82,13 +143,14 @@ function budgetLevel(remainingMs: number, totalMs: number): number { return level; } +function formatRemainingBudget(remainingMs: number): string { + if (remainingMs > FINAL_WINDOW_MS) return `about ${Math.floor(remainingMs / 60_000)} minutes`; + if (remainingMs >= 10_000) return `about ${Math.floor(remainingMs / 10_000) * 10} seconds`; + return 'under 10 seconds'; +} + function formatBudgetReminder(remainingMs: number, totalMs: number, first: boolean): string { - const remaining = - remainingMs > FINAL_WINDOW_MS - ? `about ${Math.floor(remainingMs / 60_000)} minutes` - : remainingMs >= 10_000 - ? `about ${Math.floor(remainingMs / 10_000) * 10} seconds` - : 'under 10 seconds'; + const remaining = formatRemainingBudget(remainingMs); const total = totalMs > FINAL_WINDOW_MS ? `${Math.round(totalMs / 60_000)} minutes` diff --git a/packages/local-runtime-v2/src/service/turn-system/production-composition.ts b/packages/local-runtime-v2/src/service/turn-system/production-composition.ts index fccfe3954..60311f69e 100644 --- a/packages/local-runtime-v2/src/service/turn-system/production-composition.ts +++ b/packages/local-runtime-v2/src/service/turn-system/production-composition.ts @@ -930,7 +930,12 @@ function adaptCanonicalHistory( summary: metadata.summary, ...(metadata.currentUserSourceIndex === undefined ? {} - : { currentUserSourceIndex: metadata.currentUserSourceIndex }), + : { + currentUserSourceIndex: metadata.currentUserSourceIndex, + ...(metadata.placedMarkerCount === undefined + ? {} + : { placedMarkerCount: metadata.placedMarkerCount }), + }), } as const; if (metadata.method !== 'llm_checkpoint') { const committed = await provider.compact({ diff --git a/packages/local-runtime-v2/test/integration/execution-budget-cache-prefix.integration.test.ts b/packages/local-runtime-v2/test/integration/execution-budget-cache-prefix.integration.test.ts new file mode 100644 index 000000000..8efc433a3 --- /dev/null +++ b/packages/local-runtime-v2/test/integration/execution-budget-cache-prefix.integration.test.ts @@ -0,0 +1,144 @@ +import { Type } from '@sinclair/typebox'; +import { describe, expect, it } from 'vitest'; +import { + createAssistantMessageEventStream, + type Api, + type AssistantMessage, + type Message, + type Model, +} from '@earendil-works/pi-ai'; +import type { StreamFn } from '@earendil-works/pi-agent-core'; +import { PiTurnRunner } from '@mavis/agent-core/pi-turn-runner'; +import type { RuntimeTool } from '@mavis/agent-core/tools'; + +import { createExecutionBudgetReminder } from '../../src/service/turn-system/agent-host/runner/execution-budget-reminder.js'; + +const model: Model = { + id: 'prefix-model', + name: 'prefix-model', + api: 'anthropic-messages', + provider: 'anthropic', + baseUrl: 'https://example.invalid', + reasoning: false, + input: ['text'], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 200_000, + maxTokens: 8_192, +}; + +const DEADLINE_MS = 601_000; +/** + * Clock at which each provider request of one budgeted tool loop is prepared: + * the first announcement, quiet requests, the 50% and 25% thresholds, and the + * final two minutes, including two requests inside the same 10-second step. + */ +const REQUEST_TIMES_MS = [ + 1_000, 60_000, 302_000, 320_000, 452_000, 470_000, 482_000, 485_000, 492_000, 590_000, +]; + +describe('execution budget reminder and the provider prompt cache', () => { + it('keeps every provider request an append-only extension of the previous one', async () => { + // A prompt cache entry is written at the last user message of a request. + // The next request only reads it back when that message is still in the + // same place, so the provider request must never drop or reorder earlier + // messages, including reminders, within one turn. + const clock = { nowMs: REQUEST_TIMES_MS[0]! }; + const reminder = createExecutionBudgetReminder( + DEADLINE_MS, + () => clock.nowMs, + () => true, + ); + const requests: Message[][] = []; + + await new PiTurnRunner().runTurn({ + sessionId: 'session-budget-prefix', + turnId: 'turn-budget-prefix', + workspaceDir: '/workspace', + systemPrompt: 'Use the probe tool until told to stop.', + userMessage: { text: 'run the probe loop' }, + llm: { model, streamFn: probeLoopStream(requests, clock) }, + eventWriter: { pushRuntime: async () => undefined, appendEvents: async () => undefined }, + toolConfig: { + tools: [probeTool()], + disableBuiltinToolFallback: true, + context: { sessionId: 'session-budget-prefix', turnId: 'turn-budget-prefix' }, + }, + hooks: { beforeLlmCallHook: reminder ? [reminder] : [] }, + }); + + // The loop reached its final, tool-free request. + expect(requests).toHaveLength(REQUEST_TIMES_MS.length); + requests.slice(1).forEach((request, index) => { + const previous = requests[index]!; + expect( + request.slice(0, previous.length), + `request ${index + 1} must start with request ${index}`, + ).toEqual(previous); + }); + // Every emitted reminder stays in place, and identical final-window text + // is not repeated. + expect(JSON.stringify(requests.at(-1)).match(/Execution time remaining: [^<]*/gu)).toEqual([ + 'Execution time remaining: about 10 minutes of 10 minutes total. Model generation, tools, and waiting all count against it.', + 'Execution time remaining: about 4 minutes of 10 minutes total.', + 'Execution time remaining: about 2 minutes of 10 minutes total.', + 'Execution time remaining: about 110 seconds of 10 minutes total.', + 'Execution time remaining: about 100 seconds of 10 minutes total.', + 'Execution time remaining: about 10 seconds of 10 minutes total.', + ]); + }); +}); + +function probeLoopStream(requests: Message[][], clock: { nowMs: number }): StreamFn { + return async (_model, context) => { + requests.push(structuredClone(context.messages)); + const call = requests.length; + clock.nowMs = REQUEST_TIMES_MS[call] ?? clock.nowMs; + const final = + call < REQUEST_TIMES_MS.length + ? assistant( + [{ type: 'toolCall', id: `probe-${call}`, name: 'probe', arguments: {} }], + 'toolUse', + ) + : assistant([{ type: 'text', text: 'done' }], 'stop'); + const stream = createAssistantMessageEventStream(); + stream.push({ type: 'done', reason: final.stopReason === 'toolUse' ? 'toolUse' : 'stop', message: final }); + return stream; + }; +} + +function probeTool(): RuntimeTool { + let calls = 0; + return { + def: { name: 'probe', description: 'Returns a probe result.', schema: Type.Object({}) }, + impl: { + async execute() { + calls += 1; + const text = `probe result ${calls}`; + return { tool_name: 'probe', text, content: [{ type: 'text', text }] }; + }, + }, + }; +} + +function assistant( + content: AssistantMessage['content'], + stopReason: AssistantMessage['stopReason'], +): AssistantMessage { + return { + role: 'assistant', + content, + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason, + timestamp: 0, + }; +} diff --git a/release/public-source.json b/release/public-source.json index df74562df..4e056437d 100644 --- a/release/public-source.json +++ b/release/public-source.json @@ -1406,6 +1406,7 @@ "packages/local-runtime-v2/src/service/turn-system/agent-host/production-dependencies.ts", "packages/local-runtime-v2/src/service/turn-system/agent-host/runner/contracts.ts", "packages/local-runtime-v2/src/service/turn-system/agent-host/runner/eval-reporter.ts", + "packages/local-runtime-v2/src/service/turn-system/agent-host/runner/execution-budget-reminder.test.ts", "packages/local-runtime-v2/src/service/turn-system/agent-host/runner/execution-budget-reminder.ts", "packages/local-runtime-v2/src/service/turn-system/agent-host/runner/index.ts", "packages/local-runtime-v2/src/service/turn-system/agent-host/runner/local-agent-turn-runner.ts", @@ -1530,6 +1531,7 @@ "packages/local-runtime-v2/test/helpers/plugin-database.ts", "packages/local-runtime-v2/test/integration/btw-settled-history.integration.test.ts", "packages/local-runtime-v2/test/integration/compaction-openai-transport.integration.test.ts", + "packages/local-runtime-v2/test/integration/execution-budget-cache-prefix.integration.test.ts", "packages/local-runtime-v2/test/integration/mcp/local-mcp-public-facade.integration.test.ts", "packages/local-runtime-v2/test/integration/mcp/project-mcp.integration.test.ts", "packages/local-runtime-v2/test/integration/session-stop-cascade.integration.test.ts", diff --git a/test/byok.test.mjs b/test/byok.test.mjs index 325a22f74..629e1c99b 100644 --- a/test/byok.test.mjs +++ b/test/byok.test.mjs @@ -1040,15 +1040,21 @@ function cancellationTest(cancellation) { file.endsWith("messages.jsonl"), ); assert.equal(histories.length, 1); - const roles = readFileSync(path.join(sessionsDir, histories[0]), "utf8") + const messages = readFileSync(path.join(sessionsDir, histories[0]), "utf8") .trim() .split("\n") .filter(Boolean) - .map((line) => JSON.parse(line).message?.role); + .map((line) => JSON.parse(line).message); + // `--timeout` adds the hidden execution budget reminder to history before + // the first provider request, so it survives the cancellation as well. assert.deepEqual( - roles, - cancellation === "tool" ? ["user", "assistant", "toolResult"] : ["user"], + messages.map((message) => message?.role), + cancellation === "tool" + ? ["user", "custom", "assistant", "toolResult"] + : ["user", "custom"], "Abort reconciliation must preserve executed tools without retaining cancelled text output", ); + assert.equal(messages[1].customType, "execution_budget_reminder"); + assert.equal(messages[1].display, false); }; } diff --git a/test/vitest-suites.json b/test/vitest-suites.json index 157c81564..b1f548727 100644 --- a/test/vitest-suites.json +++ b/test/vitest-suites.json @@ -119,6 +119,8 @@ "packages/tui/test/unit/runtime-error-presentation-rejection.test.ts", "packages/local-runtime-v2/src/service/turn-system/agent-host/assembly/local-turn-payload-transform.test.ts", "packages/local-runtime-v2/src/service/turn-system/agent-host/execution/executor.test.ts", + "packages/local-runtime-v2/src/service/turn-system/agent-host/runner/execution-budget-reminder.test.ts", + "packages/local-runtime-v2/test/integration/execution-budget-cache-prefix.integration.test.ts", "packages/local-runtime-v2/src/service/turn-system/execution/reminder/session-identity-reminder.test.ts", "packages/local-runtime-v2/src/services.test.ts", "packages/local-runtime-v2/src/application/session/session-title-policy.test.ts",