diff --git a/packages/agent-core-v2/docs/di-scope-domains.puml b/packages/agent-core-v2/docs/di-scope-domains.puml index 12c8da502..de989f2af 100644 --- a/packages/agent-core-v2/docs/di-scope-domains.puml +++ b/packages/agent-core-v2/docs/di-scope-domains.puml @@ -301,6 +301,7 @@ fullCompaction --> contextMemory #34495E fullCompaction --> contextSize #34495E fullCompaction --> llmRequester #34495E fullCompaction --> profile #34495E +fullCompaction --> toolRegistry #34495E fullCompaction --> telemetry #34495E fullCompaction --> wireRecord #34495E fullCompaction --> turn #34495E diff --git a/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts b/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts index da9372a52..6f670e49a 100644 --- a/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts +++ b/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts @@ -2,7 +2,12 @@ import { Disposable } from "#/_base/di/lifecycle"; import { InstantiationType } from '#/_base/di/extensions'; import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; import { renderPrompt } from "#/_base/utils/render-prompt"; -import { estimateTokensForMessage, estimateTokensForMessages } from "#/_base/utils/tokens"; +import { + estimateTokens, + estimateTokensForMessage, + estimateTokensForMessages, + estimateTokensForTools, +} from "#/_base/utils/tokens"; import { buildCompactionSummaryText, isRealUserInput } from '#/agent/contextMemory/compactionHandoff'; import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; import type { ContextMessage } from '#/agent/contextMemory/types'; @@ -11,7 +16,8 @@ import { IAgentLLMRequesterService, type LLMRequestFinish } from '#/agent/llmReq import { retryBackoffDelays, sleepForRetry } from '#/agent/llmRequester/retry'; import { IAgentLoopService, type LoopErrorContext } from '#/agent/loop/loop'; import { isAbortError, isContextOverflowError } from '#/agent/loop/errors'; -import { IAgentProfileService } from '#/agent/profile/profile'; +import { IAgentProfileService, type ProfileModelContext } from '#/agent/profile/profile'; +import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry'; import { IAgentTurnService } from '#/agent/turn/turn'; import { ISessionTodoService } from '#/session/todo/sessionTodo'; import { renderTodoList, type TodoItem } from '#/session/todo/todoItem'; @@ -22,6 +28,7 @@ import { isRetryableGenerateError, } from '#/app/llmProtocol/errors'; import { createUserMessage, type Message } from '#/app/llmProtocol/message'; +import type { Tool } from '#/app/llmProtocol/tool'; import { type TokenUsage } from '#/app/llmProtocol/usage'; import { IEventBus } from '#/app/event/eventBus'; import { ITelemetryService } from '#/app/telemetry/telemetry'; @@ -67,9 +74,14 @@ declare module '#/agent/wireRecord/wireRecord' { export const MAX_COMPACTION_RETRY_ATTEMPTS = 5; const DEFAULT_COMPACTION_MAX_COMPLETION_TOKENS = 128 * 1024; +const OVERFLOW_CONTEXT_SAFETY_RATIO = 0.85; const OVERFLOW_STATUS_RECOVERY_RATIO = 0.5; const MAX_COMPACTION_OVERFLOW_SHRINK_ATTEMPTS = 3; const COMPACTION_OVERFLOW_SHRINK_RATIOS = [0.7, 0.5, 0.35] as const; +const EMPTY_TOOL_PARAMETERS: Record = { + type: 'object', + properties: {}, +}; type CompactionTelemetryProperties = Record; @@ -98,6 +110,7 @@ export class AgentFullCompactionService extends Disposable implements IAgentFull private readonly strategy: CompactionStrategy; private compactionCountInTurn = 0; private _compacting: ActiveCompaction | null = null; + private readonly observedMaxContextTokensByModel = new Map(); // Token count right after the last successful compaction. While nothing new // has been appended, the history is already in its minimal compacted form; // re-compacting would only summarize the summary again, so @@ -114,6 +127,7 @@ export class AgentFullCompactionService extends Disposable implements IAgentFull @IAgentContextSizeService private readonly contextSize: IAgentContextSizeService, @IAgentLLMRequesterService private readonly llmRequester: IAgentLLMRequesterService, @IAgentProfileService private readonly profile: IAgentProfileService, + @IAgentToolRegistryService private readonly toolRegistry: IAgentToolRegistryService, @ISessionTodoService private readonly todo: ISessionTodoService, @ITelemetryService private readonly telemetry: ITelemetryService, @IAgentWireService private readonly wire: IWireService, @@ -122,7 +136,7 @@ export class AgentFullCompactionService extends Disposable implements IAgentFull @IAgentLoopService loopService: IAgentLoopService, ) { super(); - this.strategy = new RuntimeCompactionStrategy(() => this.profile.resolveModelContext()); + this.strategy = new RuntimeCompactionStrategy(() => this.resolveModelContextWithEffectiveMax()); this._register(this.wire.onRestored(() => this.normalizeAfterReplay())); this._register( this.eventBus.subscribe('turn.started', () => this.resetForTurn()), @@ -150,6 +164,78 @@ export class AgentFullCompactionService extends Disposable implements IAgentFull return this._compacting; } + private getEffectiveMaxContextTokens(): number { + const configured = this.profile.data().modelCapabilities.max_context_tokens; + const modelAlias = this.profile.data().modelAlias; + const observed = + modelAlias === undefined ? undefined : this.observedMaxContextTokensByModel.get(modelAlias); + if (observed === undefined) return configured; + if (configured <= 0) return observed; + return Math.min(configured, observed); + } + + private resolveModelContextWithEffectiveMax(): ProfileModelContext { + const resolved = this.profile.resolveModelContext(); + return { + ...resolved, + modelCapabilities: { + ...resolved.modelCapabilities, + max_context_tokens: this.getEffectiveMaxContextTokens(), + }, + }; + } + + private estimateCurrentRequestTokens(): number { + return this.estimateRequestTokens(this.context.get()); + } + + private estimateRequestTokens(messages: readonly Message[]): number { + return ( + estimateTokens(this.profile.getSystemPrompt()) + + estimateTokensForTools(this.defaultTools().filter((tool) => tool.deferred !== true)) + + estimateTokensForMessages(messages) + ); + } + + private defaultTools(): readonly Tool[] { + return this.toolRegistry + .list() + .filter((tool) => this.profile.isToolActive(tool.name, tool.source)) + .map((tool) => ({ + name: tool.name, + description: tool.description, + parameters: tool.parameters ?? EMPTY_TOOL_PARAMETERS, + })); + } + + private shouldRecoverFromContextOverflow( + error: unknown, + estimatedRequestTokens = this.estimateCurrentRequestTokens(), + ): boolean { + if (isContextOverflowError(error)) return true; + const statusError = findAPIStatusError(error); + if (statusError instanceof APIContextOverflowError) return true; + if (statusError === undefined || statusError.statusCode !== 413) return false; + const effectiveMax = this.getEffectiveMaxContextTokens(); + return ( + effectiveMax > 0 && + estimatedRequestTokens >= effectiveMax * OVERFLOW_STATUS_RECOVERY_RATIO + ); + } + + private observeContextOverflow(estimatedRequestTokens: number): void { + if (!Number.isFinite(estimatedRequestTokens) || estimatedRequestTokens <= 0) return; + const modelAlias = this.profile.data().modelAlias; + if (modelAlias === undefined) return; + const observed = Math.max( + 1, + Math.floor(estimatedRequestTokens * OVERFLOW_CONTEXT_SAFETY_RATIO), + ); + const current = this.getEffectiveMaxContextTokens(); + if (current > 0 && observed >= current) return; + this.observedMaxContextTokensByModel.set(modelAlias, observed); + } + begin(input: FullCompactionInput): boolean { if (this._compacting) return false; const data: CompactionBeginData = { source: input.source, instruction: input.instruction }; @@ -238,14 +324,16 @@ export class AgentFullCompactionService extends Disposable implements IAgentFull context: LoopErrorContext, next: () => Promise, ): Promise { - const isOverflow = - isContextOverflowError(context.error) || - findAPIStatusError(context.error) instanceof APIContextOverflowError || - this.shouldRecoverFromPlain413(context.error); + const estimatedRequestTokens = this.estimateCurrentRequestTokens(); + const isOverflow = this.shouldRecoverFromContextOverflow( + context.error, + estimatedRequestTokens, + ); if (!isOverflow) { await next(); return; } + this.observeContextOverflow(estimatedRequestTokens); this.consecutiveOverflowCompactions += 1; const maxAttempts = this.strategy.maxOverflowCompactionAttempts; if (this.consecutiveOverflowCompactions > maxAttempts) { @@ -400,6 +488,7 @@ export class AgentFullCompactionService extends Disposable implements IAgentFull // Raw context slice — `llmRequester` projects every request once; // projecting here too would run micro-compaction on shifted indices. const messages: Message[] = [...messagesToCompact, createUserMessage(instruction)]; + const estimatedCompactionRequestTokens = this.estimateRequestTokens(messages); try { attempt = collectSummary( @@ -418,7 +507,12 @@ export class AgentFullCompactionService extends Disposable implements IAgentFull ); break; } catch (error) { - if (this.shouldRecoverFromCompactionOverflow(error, messages)) { + const isContextOverflow = this.shouldRecoverFromContextOverflow( + error, + estimatedCompactionRequestTokens, + ); + if (isContextOverflow) { + this.observeContextOverflow(estimatedCompactionRequestTokens); overflowShrinkCount += 1; if ( overflowShrinkCount > MAX_COMPACTION_OVERFLOW_SHRINK_ATTEMPTS || @@ -528,27 +622,6 @@ export class AgentFullCompactionService extends Disposable implements IAgentFull private tokenCountWithPending(): number { return this.contextSize.get().size; } - - private shouldRecoverFromPlain413( - error: unknown, - estimatedRequestTokens = this.tokenCountWithPending(), - ): boolean { - const statusError = findAPIStatusError(error); - if (statusError === undefined || statusError.statusCode !== 413) return false; - const maxContextTokens = this.profile.getModelCapabilities().max_context_tokens; - return ( - maxContextTokens > 0 && - estimatedRequestTokens >= maxContextTokens * OVERFLOW_STATUS_RECOVERY_RATIO - ); - } - - private shouldRecoverFromCompactionOverflow( - error: unknown, - messages: readonly Message[], - ): boolean { - if (error instanceof APIContextOverflowError) return true; - return this.shouldRecoverFromPlain413(error, estimateTokensForMessages(messages)); - } } function findAPIStatusError(error: unknown): APIStatusError | undefined { diff --git a/packages/agent-core-v2/test/fullCompaction/full.test.ts b/packages/agent-core-v2/test/fullCompaction/full.test.ts index d842e82eb..9511812cc 100644 --- a/packages/agent-core-v2/test/fullCompaction/full.test.ts +++ b/packages/agent-core-v2/test/fullCompaction/full.test.ts @@ -1845,6 +1845,79 @@ describe('FullCompaction', () => { await ctx.expectResumeMatches(); }); + it('remembers the observed provider context window after overflow', async () => { + let callCount = 0; + const generate: GenerateFn = async (_provider, _system, _tools, _history, callbacks) => { + callCount += 1; + if (callCount === 1) { + throw new APIContextOverflowError(400, 'Context length exceeded', 'req-observed-window'); + } + if (callCount === 2) { + return textResult('Observed recovery summary.'); + } + if (callCount === 3) { + await callbacks?.onMessagePart?.({ + type: 'text', + text: 'Recovered after observed overflow.', + }); + return textResult('Recovered after observed overflow.'); + } + if (callCount === 4) { + return textResult('Observed preemptive summary.'); + } + if (callCount === 5) { + await callbacks?.onMessagePart?.({ + type: 'text', + text: 'Answered after observed-window precompaction.', + }); + return textResult('Answered after observed-window precompaction.'); + } + throw new Error(`Unexpected generate call ${String(callCount)}`); + }; + const ctx = testAgent({ generate }); + ctx.configure({ + provider: CATALOGUED_PROVIDER, + modelCapabilities: { + ...CATALOGUED_MODEL_CAPABILITIES, + max_context_tokens: 200_000, + }, + tools: SNAPSHOT_VISIBLE_TOOLS, + }); + ctx.appendExchange(1, 'old user one', 'old assistant one', 20); + ctx.newEvents(); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'learn observed window' }] }); + await ctx.untilTurnEnd(); + expect(callCount).toBe(3); + + ctx.appendExchange(2, 'near observed user', 'near observed assistant', 120_000); + ctx.newEvents(); + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'use observed window' }] }); + const events = await ctx.untilTurnEnd(); + + expect(callCount).toBe(5); + expect(eventIndex(events, 'compaction.started')).toBeLessThan( + eventIndex(events, 'turn.step.started'), + ); + expect(events).toContainEqual( + expect.objectContaining({ + event: 'compaction.completed', + args: expect.objectContaining({ + result: expect.objectContaining({ + summary: 'Observed preemptive summary.', + }), + }), + }), + ); + expect(events).toContainEqual( + expect.objectContaining({ + event: 'turn.ended', + args: { turnId: 1, reason: 'completed' }, + }), + ); + await ctx.expectResumeMatches(); + }); + it('recovers from plain 413 when estimated request is over effective max', async () => { let callCount = 0; const generate: GenerateFn = async (_provider, _system, _tools, _history, callbacks) => {