diff --git a/packages/runtime/src/__tests__/ai-sdk-backend.test.ts b/packages/runtime/src/__tests__/ai-sdk-backend.test.ts index f9693fd817..ade55a6edc 100644 --- a/packages/runtime/src/__tests__/ai-sdk-backend.test.ts +++ b/packages/runtime/src/__tests__/ai-sdk-backend.test.ts @@ -4527,6 +4527,92 @@ describe('AiSdkBackend model history', () => { assert.equal(result.contextBudget?.compactionDecisions?.[0]?.decision, 'replaced'); }); + test('coalesces identical in-flight compactHistory requests', async () => { + const recorded: HistoryCompactCheckpoint[] = []; + let summarizeCalls = 0; + let releaseSummary!: () => void; + const summaryReady = new Promise((resolve) => { + releaseSummary = resolve; + }); + const backend = createTestAiSdkBackend({ + sessionId: 'session-1', + header: header(), + appendMessage: async () => {}, + connection: connection(), + apiKey: 'sk-test', + modelId: 'mock-model-id', + modelFactory: () => completionModel(), + tools: [], + newId: idGenerator(), + now: monotonicClock(), + contextBudget: { + name: 'in-flight-dedup-test', + charsPerToken: 1, + }, + summarizeHistoryCompact: async () => { + summarizeCalls += 1; + await summaryReady; + return structuredSummary('IN_FLIGHT_DEDUP_SUMMARY'); + }, + recordHistoryCompactCheckpoint: (checkpoint) => { + recorded.push(checkpoint); + }, + }); + const runtimeContext = [ + runtimeTextEvent({ + id: 'dedup-old-user', + turnId: 'dedup-old-turn', + role: 'user', + author: 'user', + text: 'old context '.repeat(100), + }), + runtimeTextEvent({ + id: 'dedup-old-agent', + turnId: 'dedup-old-turn', + role: 'model', + author: 'agent', + text: 'old response '.repeat(100), + }), + ]; + const first = backend.compactHistory({ + turnId: 'dedup-compact-1', + runId: 'run-dedup-1', + runtimeContext, + }); + await new Promise((resolve) => setImmediate(resolve)); + const second = backend.compactHistory({ + turnId: 'dedup-compact-2', + runId: 'run-dedup-2', + runtimeContext: [ + ...runtimeContext, + runtimeTextEvent({ + id: 'dedup-current-turn', + turnId: 'dedup-compact-2', + role: 'user', + author: 'user', + text: 'current turn content must not affect the fold key', + }), + ], + }); + + releaseSummary(); + const [firstResult, secondResult] = await Promise.all([first, second]); + assert.equal(summarizeCalls, 1); + assert.equal(secondResult.outcome.kind, 'compacted'); + assert.equal(firstResult.outcome.kind, 'compacted'); + assert.equal(recorded.length, 2); + const firstCheckpoint = recorded[0]; + const secondCheckpoint = recorded[1]; + assert.ok(firstCheckpoint); + assert.ok(secondCheckpoint); + assert.equal(firstCheckpoint.version, 2); + assert.equal(secondCheckpoint.version, 2); + if (firstCheckpoint.version === 2 && secondCheckpoint.version === 2) { + assert.equal(secondCheckpoint.summary, firstCheckpoint.summary); + assert.deepEqual(secondCheckpoint.coverage, firstCheckpoint.coverage); + } + }); + test('manual compactHistory compacts one completed turn with multiple agent steps', async () => { const recorded: HistoryCompactCheckpoint[] = []; const backend = createTestAiSdkBackend({ diff --git a/packages/runtime/src/ai-sdk-compaction.ts b/packages/runtime/src/ai-sdk-compaction.ts index 4199e4541a..56ee552246 100644 --- a/packages/runtime/src/ai-sdk-compaction.ts +++ b/packages/runtime/src/ai-sdk-compaction.ts @@ -57,6 +57,7 @@ import { matchHistoryCompactCheckpointPrefix, projectHistoryCompactCheckpointReplay, type HistoryCompactCheckpoint, + type HistoryCompactCheckpointHeadAnchor, type HistoryCompactMemoryExtractionBoundary, type HistoryCompactProviderState, } from './history-compact-checkpoint.js'; @@ -154,6 +155,12 @@ export interface AutomaticMemoryCompactionDecision { readonly dispatch: boolean; } +interface InFlightHistorySummary { + readonly abortController: AbortController; + readonly consumers: Set; + readonly promise: Promise; +} + /** Constructor dependencies for AiSdkCompaction. */ export interface AiSdkCompactionDeps { input: AiSdkCompactionCapabilities; @@ -209,7 +216,14 @@ export class AiSdkCompaction { providerReasoningReplayEventIds: ReadonlySet, ) => Promise; private readonly canReplayProviderNative: (plan: RuntimeEventModelReplayPlan) => boolean; - private historyCompactAbortController: AbortController | null = null; + private readonly historyCompactAbortControllers = new Set(); + /** + * Exact duplicate summary inputs share one physical provider call. The map + * lives at the summarizer boundary, where the existing effective-history + * fingerprint already describes the request that spends provider budget. + * Each caller still completes its own checkpoint/turn bookkeeping. + */ + private readonly inFlightHistorySummaries = new Map(); /** * Session-scoped circuit for exact malformed compaction inputs. A retry or * regeneration on the same backend must not dispatch the same doomed call; @@ -281,7 +295,7 @@ export class AiSdkCompaction { /** Abort an in-flight manual history compaction (called by AiSdkBackend.stop). */ public abortHistoryCompact(): void { - this.historyCompactAbortController?.abort(); + for (const controller of this.historyCompactAbortControllers) controller.abort(); } public async compactHistory( @@ -289,7 +303,7 @@ export class AiSdkCompaction { automaticMemoryBoundary?: HistoryCompactMemoryExtractionBoundary, ): Promise { const historyCompactAbortController = new AbortController(); - this.historyCompactAbortController = historyCompactAbortController; + this.historyCompactAbortControllers.add(historyCompactAbortController); try { const policy = this.input.contextBudget; const summarizer = this.input.summarizeHistoryCompact; @@ -384,21 +398,35 @@ export class AiSdkCompaction { ...(automaticMemoryBoundary ? { memoryExtractionBoundary: automaticMemoryBoundary } : {}), ...(previousCheckpoint ? { previousCheckpoint } : {}), summarize: async ({ coveredRuntimeEvents, newlyFoldedRuntimeEvents, previousCheckpoint }) => - await this.summarizeWithFailureCircuit(summarizer, { - sessionId: this.sessionId, - turnId: input.turnId, - runId: input.runId, - source: { - foldedRuntimeEvents: [...coveredRuntimeEvents], - ...(input.runtimeContextInvocations - ? { invocations: input.runtimeContextInvocations } + await this.summarizeWithFailureCircuit( + summarizer, + { + sessionId: this.sessionId, + turnId: input.turnId, + runId: input.runId, + source: { + foldedRuntimeEvents: [...coveredRuntimeEvents], + ...(input.runtimeContextInvocations + ? { + invocations: invocationsForFoldedEvents( + input.runtimeContextInvocations, + coveredRuntimeEvents, + ), + } + : {}), + }, + newlyFoldedRuntimeEvents: [...newlyFoldedRuntimeEvents], + ...(previousCheckpoint ? { previousCheckpoint } : {}), + abortSignal: historyCompactAbortController.signal, + ...(tracker ? { providerRequestTracker: tracker } : {}), + }, + { + phase: 'pre_turn', + ...(automaticMemoryBoundary + ? { memoryExtractionBoundary: automaticMemoryBoundary } : {}), }, - newlyFoldedRuntimeEvents: [...newlyFoldedRuntimeEvents], - ...(previousCheckpoint ? { previousCheckpoint } : {}), - abortSignal: historyCompactAbortController.signal, - ...(tracker ? { providerRequestTracker: tracker } : {}), - }), + ), }); if (historyCompactAbortController.signal.aborted) { return { outcome: { kind: 'failed', reason: 'aborted' } }; @@ -466,9 +494,7 @@ export class AiSdkCompaction { }), }; } finally { - if (this.historyCompactAbortController === historyCompactAbortController) { - this.historyCompactAbortController = null; - } + this.historyCompactAbortControllers.delete(historyCompactAbortController); } } @@ -479,11 +505,14 @@ export class AiSdkCompaction { private async summarizeWithFailureCircuit( summarizer: HistoryCompactSummarizer, input: HistoryCompactSummaryInput, + checkpointIntent?: { + phase?: 'pre_turn' | 'mid_turn'; + headAnchor?: HistoryCompactCheckpointHeadAnchor; + memoryExtractionBoundary?: HistoryCompactMemoryExtractionBoundary; + }, ): Promise { - const foldedRunIds = new Set(input.source.foldedRuntimeEvents.map((event) => event.runId)); const sourceRunRoutes = input.source.invocations - ?.filter((invocation) => foldedRunIds.has(invocation.runId)) - .map((invocation) => { + ?.map((invocation) => { const route = invocation.opening.route; return { runId: invocation.runId, @@ -509,13 +538,32 @@ export class AiSdkCompaction { sourceRunRoutes, foldedRuntimeEvents: input.source.foldedRuntimeEvents, newlyFoldedRuntimeEvents: input.newlyFoldedRuntimeEvents, + checkpointIntent, }), ); const priorFailure = this.malformedSummaryFailures.get(fingerprint); if (priorFailure) throw new HistoryCompactSummarizerError(priorFailure); + let shared = this.inFlightHistorySummaries.get(fingerprint); + if (!shared) { + const abortController = new AbortController(); + const pending = Promise.resolve().then(() => + summarizer({ ...input, abortSignal: abortController.signal }), + ); + shared = { abortController, consumers: new Set(), promise: pending }; + this.inFlightHistorySummaries.set(fingerprint, shared); + void pending.then( + () => this.removeInFlightHistorySummary(fingerprint, shared!), + () => this.removeInFlightHistorySummary(fingerprint, shared!), + ); + } + const consumer = Symbol('history-summary-consumer'); + shared.consumers.add(consumer); try { - return await Promise.resolve(summarizer(input)); + // A rider must be able to stop waiting without aborting the physical + // call for other consumers. The shared controller is aborted only when + // every consumer has detached. + return await waitForAbortablePromise(shared.promise, input.abortSignal); } catch (error) { if ( error instanceof HistoryCompactSummarizerError && @@ -530,6 +578,20 @@ export class AiSdkCompaction { } } throw error; + } finally { + shared.consumers.delete(consumer); + if ( + shared.consumers.size === 0 && + this.inFlightHistorySummaries.get(fingerprint) === shared + ) { + shared.abortController.abort(); + } + } + } + + private removeInFlightHistorySummary(fingerprint: string, shared: InFlightHistorySummary): void { + if (this.inFlightHistorySummaries.get(fingerprint) === shared) { + this.inFlightHistorySummaries.delete(fingerprint); } } @@ -1077,6 +1139,15 @@ export class AiSdkCompaction { } const orderedEvents = [...state.priorContentEvents, ...currentTurnEvents]; const memoryDecision = input.memoryCompactionDecision?.(); + const memoryExtractionBoundary = + memoryDecision && orderedEvents.at(-1) + ? { + runId: orderedEvents.at(-1)!.runId, + turnId: orderedEvents.at(-1)!.turnId, + runtimeEventId: orderedEvents.at(-1)!.id, + disposition: memoryDecision.disposition, + } + : undefined; const plan = await planHistoryCompaction({ sessionId: this.sessionId, phase: input.phase ?? 'mid_turn', @@ -1094,30 +1165,29 @@ export class AiSdkCompaction { ? { highWaterName: compactPolicy.highWaterName } : {}), ...(state.previousCheckpoint ? { previousCheckpoint: state.previousCheckpoint } : {}), - ...(memoryDecision && orderedEvents.at(-1) - ? { - memoryExtractionBoundary: { - runId: orderedEvents.at(-1)!.runId, - turnId: orderedEvents.at(-1)!.turnId, - runtimeEventId: orderedEvents.at(-1)!.id, - disposition: memoryDecision.disposition, - }, - } - : {}), + ...(memoryExtractionBoundary ? { memoryExtractionBoundary } : {}), summarize: async ({ coveredRuntimeEvents, newlyFoldedRuntimeEvents, previousCheckpoint }) => { - return await this.summarizeWithFailureCircuit(summarizer, { - sessionId: this.sessionId, - turnId, - ...(input.origin.runId ? { runId: input.origin.runId } : {}), - source: { - foldedRuntimeEvents: [...coveredRuntimeEvents], - invocations: state.priorInvocations, + return await this.summarizeWithFailureCircuit( + summarizer, + { + sessionId: this.sessionId, + turnId, + ...(input.origin.runId ? { runId: input.origin.runId } : {}), + source: { + foldedRuntimeEvents: [...coveredRuntimeEvents], + invocations: invocationsForFoldedEvents(state.priorInvocations, coveredRuntimeEvents), + }, + ...(previousCheckpoint ? { previousCheckpoint } : {}), + newlyFoldedRuntimeEvents: [...newlyFoldedRuntimeEvents], + ...(abortSignal ? { abortSignal } : {}), + ...(midTurnTracker ? { providerRequestTracker: midTurnTracker } : {}), }, - ...(previousCheckpoint ? { previousCheckpoint } : {}), - newlyFoldedRuntimeEvents: [...newlyFoldedRuntimeEvents], - ...(abortSignal ? { abortSignal } : {}), - ...(midTurnTracker ? { providerRequestTracker: midTurnTracker } : {}), - }); + { + phase: input.phase ?? 'mid_turn', + headAnchor: { runtimeEventId: state.headAnchor.id, turnId }, + ...(memoryExtractionBoundary ? { memoryExtractionBoundary } : {}), + }, + ); }, }); @@ -1634,6 +1704,48 @@ function waitForQueueProgressOrAbort( }); } +/** Let an individual compaction stop waiting without cancelling shared work. */ +function waitForAbortablePromise( + promise: Promise, + abortSignal?: AbortSignal, +): Promise { + if (!abortSignal) return promise; + if (abortSignal.aborted) return Promise.resolve(undefined); + return new Promise((resolve, reject) => { + let settled = false; + const cleanup = () => abortSignal.removeEventListener('abort', onAbort); + const onAbort = () => { + if (settled) return; + settled = true; + cleanup(); + resolve(undefined); + }; + abortSignal.addEventListener('abort', onAbort, { once: true }); + void promise.then( + (value) => { + if (settled) return; + settled = true; + cleanup(); + resolve(value); + }, + (error: unknown) => { + if (settled) return; + settled = true; + cleanup(); + reject(error); + }, + ); + }); +} + +function invocationsForFoldedEvents( + invocations: readonly RuntimeInvocationRecord[], + foldedRuntimeEvents: readonly RuntimeEvent[], +): RuntimeInvocationRecord[] { + const foldedRunIds = new Set(foldedRuntimeEvents.map((event) => event.runId)); + return invocations.filter((invocation) => foldedRunIds.has(invocation.runId)); +} + export function hasBlockingReplayDiagnostics(plan: RuntimeEventModelReplayPlan): boolean { // `unmatched_tool_result` is deliberately NOT blocking: the materializer // drops an orphan tool result (its call sliced away or the ledger corrupt) diff --git a/packages/runtime/src/history-compact-checkpoint-coordinator.ts b/packages/runtime/src/history-compact-checkpoint-coordinator.ts index 05533e500f..3ee0a9b85a 100644 --- a/packages/runtime/src/history-compact-checkpoint-coordinator.ts +++ b/packages/runtime/src/history-compact-checkpoint-coordinator.ts @@ -90,12 +90,23 @@ export class HistoryCompactCheckpointCoordinator { .catch(() => {}) .then(async () => { const durableCheckpoint = await this.load(sessionId); - if (!canReplaceHistoryCompactCheckpoint(durableCheckpoint, checkpoint)) { + const sameEffectiveCheckpoint = hasSameEffectiveCoverage(durableCheckpoint, checkpoint); + if ( + !sameEffectiveCheckpoint && + !canReplaceHistoryCompactCheckpoint(durableCheckpoint, checkpoint) + ) { throw new Error('History compact checkpoint was superseded before persistence'); } - await run.recordHistoryCompactCheckpoint(checkpoint); - this.checkpoints.set(sessionId, checkpoint); - this.scheduleCleanup(sessionId, checkpoint); + // A concurrent caller may have received the same effective checkpoint + // from the summarizer coalescer. Keep its run-local ledger event too so + // the rider Turn retains provenance even though the session checkpoint + // itself is already current. + const checkpointToRecord = sameEffectiveCheckpoint ? durableCheckpoint! : checkpoint; + await run.recordHistoryCompactCheckpoint(checkpointToRecord); + if (!sameEffectiveCheckpoint) { + this.checkpoints.set(sessionId, checkpoint); + this.scheduleCleanup(sessionId, checkpoint); + } }) .finally(() => { if (this.writes.get(sessionId) === tracked) { @@ -155,3 +166,21 @@ export class HistoryCompactCheckpointCoordinator { this.cleanups.set(sessionId, tracked); } } + +function hasSameEffectiveCoverage( + current: HistoryCompactCheckpoint | undefined, + candidate: HistoryCompactCheckpoint, +): boolean { + if (!current) return false; + const stable = (checkpoint: HistoryCompactCheckpoint): string => { + const { + checkpointId: _checkpointId, + createdAt: _createdAt, + highWaterSeq: _highWaterSeq, + previousCheckpointId: _previousCheckpointId, + ...rest + } = checkpoint; + return JSON.stringify(rest); + }; + return stable(current) === stable(candidate); +}