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",