Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/compaction-model-option.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@moonshot-ai/agent-core-v2": minor
---

Add an opt-in dedicated compaction model: when the `compaction-model` experiment is enabled and `[compaction_model]` is configured, context compaction uses that model instead of the current one, and transparently falls back to the current model if it errors or is inaccessible.
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,13 @@ import { TurnStarted } from '#/agent/loop/turnEvents';
import { TurnEnded } from '#/agent/loop/turnOps';
import { isAbortError } from '#/_base/utils/abort';
import { IAgentProfileService, type ProfileModelContext } from '#/agent/profile/profile';
import { IConfigService } from '#/app/config/config';
import { IFlagService } from '#/app/flag/flag';
import {
compactionDisplayModel,
compactionModelBindingFor,
wrapCompactionModelError,
} from '#/session/compaction/configSection';
import {
agentContextOfScope,
IAgentScopeContext,
Expand Down Expand Up @@ -152,6 +159,8 @@ export class AgentFullCompactionService extends Service implements IAgentFullCom
@ILogService private readonly log: ILogService,
@IAgentLoopService private readonly loopService: IAgentLoopService,
@IAgentStateService private readonly states: IAgentStateService,
@IConfigService private readonly configService: IConfigService,
@IFlagService private readonly flags: IFlagService,
) {
super();
this.states.contributeState(fullCompactionKey);
Expand Down Expand Up @@ -634,12 +643,41 @@ export class AgentFullCompactionService extends Service implements IAgentFullCom
const resolvedModel = this.profile.resolveModelContext();
thinkingEffort = resolvedModel.thinkingLevel;
const maxContextTokens = resolvedModel.modelCapabilities.max_context_tokens;
const currentModelAlias = resolvedModel.modelAlias;
const defaultCompactionCap =
maxContextTokens > 0
? Math.min(maxContextTokens, DEFAULT_COMPACTION_MAX_COMPLETION_TOKENS)
: undefined;
const compactionMaxOutputSize = resolvedModel.maxOutputSize ?? defaultCompactionCap;

const binding = compactionModelBindingFor(this.configService, this.flags, {
modelAlias: currentModelAlias,
thinkingLevel: thinkingEffort,
});
const dedicatedModelAlias = binding.model;
let hasDedicatedModel = dedicatedModelAlias !== currentModelAlias;
let usingFallbackModel = false;
let boundModel = resolvedModel;
if (hasDedicatedModel) {
try {
boundModel = this.profile.resolveModelContextFor(dedicatedModelAlias);
} catch (error) {
this.log.warn(
`compaction model "${dedicatedModelAlias}" is not configured; falling back to current model "${currentModelAlias}"`,
{ cause: wrapCompactionModelError(error, dedicatedModelAlias) },
);
hasDedicatedModel = false;
boundModel = resolvedModel;
}
}
const boundMaxOutputSize = boundModel.maxOutputSize ?? defaultCompactionCap;
let compactionRequestModel = hasDedicatedModel ? dedicatedModelAlias : undefined;
let effectiveModelAlias = hasDedicatedModel ? dedicatedModelAlias : currentModelAlias;
const effectiveMaxOutputSize = Math.min(
compactionMaxOutputSize ?? DEFAULT_COMPACTION_MAX_COMPLETION_TOKENS,
boundMaxOutputSize ?? DEFAULT_COMPACTION_MAX_COMPLETION_TOKENS,
);

const customInstruction = data.instruction?.trim() ?? '';
const instruction = renderPrompt(compactionInstructionTemplate, {
custom_instruction_block:
Expand All @@ -661,7 +699,8 @@ export class AgentFullCompactionService extends Service implements IAgentFullCom
const request = this.llmRequester.start(
{
messages,
maxOutputSize: compactionMaxOutputSize,
maxOutputSize: effectiveMaxOutputSize,
model: compactionRequestModel,
source: {
type: 'operation',
turnId: active.originTurnId,
Expand Down Expand Up @@ -716,6 +755,22 @@ export class AgentFullCompactionService extends Service implements IAgentFullCom
retryCount = 0;
continue;
}
if (
hasDedicatedModel &&
!usingFallbackModel &&
(isRetryableGenerateError(unwrappedError) ||
!(error instanceof CompactionTruncatedError))
) {
usingFallbackModel = true;
effectiveModelAlias = currentModelAlias;
compactionRequestModel = undefined;
this.log.warn(
`compaction model "${dedicatedModelAlias}" failed; falling back to current model "${currentModelAlias}"`,
{ cause: wrapCompactionModelError(error, dedicatedModelAlias) },
);
retryCount = 0;
continue;
}
if (!isRetryableGenerateError(unwrappedError)) {
throw error;
}
Expand Down Expand Up @@ -764,6 +819,8 @@ export class AgentFullCompactionService extends Service implements IAgentFullCom
round: 1,
thinking_effort: thinkingEffort,
trace_id: attempt.traceId,
model: effectiveModelAlias,
model_display: compactionDisplayModel(this.configService, effectiveModelAlias),
...usageTelemetry(attempt.usage),
};
this.telemetry.track2('compaction_finished', properties);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ export interface AgentLLMRequestOverrides {
systemPrompt?: string;
source?: AgentLLMRequestSource;
maxOutputSize?: number;
model?: string;
}

export interface AgentLLMRequestTask {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -582,7 +582,10 @@ export class AgentLLMRequesterService implements IAgentLLMRequesterService {

private resolveRequest(overrides: AgentLLMRequestOverrides): ResolvedLLMRequest {
const turnConfig = this.resolveTurnConfig(overrides.source);
const resolved = turnConfig?.resolved ?? this.profile.resolveModelContext();
const resolved =
overrides.model !== undefined
? this.profile.resolveModelContextFor(overrides.model)
: turnConfig?.resolved ?? this.profile.resolveModelContext();
const baseParams = turnConfig?.params ?? this.profile.resolveRequestParams();
const budgetParams = completionBudgetParams({
budget: resolveCompletionBudget({
Expand Down
1 change: 1 addition & 0 deletions packages/agent-core-v2/src/agent/profile/profile.ts
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,7 @@ export interface IAgentProfileService {
data(): ProfileData;
getEffectiveThinkingLevel(): ThinkingEffort;
resolveModelContext(): ProfileModelContext;
resolveModelContextFor(modelAlias: string): ProfileModelContext;
resolveRequestParams(): ModelRequestParams;
getModelCapabilities(): ModelCapability;
getMaxOutputSize(): number | undefined;
Expand Down
14 changes: 14 additions & 0 deletions packages/agent-core-v2/src/agent/profile/profileService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -504,6 +504,20 @@ export class AgentProfileService extends Disposable implements IAgentProfileServ
};
}

resolveModelContextFor(modelAlias: string): ProfileModelContext {
const model = this.modelCatalog.get(modelAlias);
const loopControl = this.config.get<LoopControl>('loopControl');
return {
modelAlias,
modelCapabilities: model.capabilities,
maxOutputSize: model.maxOutputSize,
alwaysThinking: model.alwaysThinking || undefined,
thinkingLevel: this.resolveThinkingState(model).effective,
reservedContextSize: loopControl?.reservedContextSize,
compactionTriggerRatio: loopControl?.compactionTriggerRatio,
};
}

resolveRequestParams(): ModelRequestParams {
const model = this.tryResolveRawModel();
const thinking = this.resolveThinkingState(model);
Expand Down
104 changes: 104 additions & 0 deletions packages/agent-core-v2/src/app/kosongConfig/compactionModelOverlay.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
import type { ConfigEffectiveOverlay } from '#/app/config/config';
import { registerConfigOverlay } from '#/app/config/configOverlayContributions';
import { isPlainObject } from '#/app/config/toml';
import type { ModelOverride } from '#/kosong/model/model';

import {
COMPACTION_MODEL_SECTION,
DEFAULT_MODEL_SECTION,
MODELS_SECTION,
type CompactionModelConfig,
} from './configSection';

/**
* `kosongConfig` domain — `[compaction_model]` derived-entry overlay.
*
* Compaction-model mirror of {@link visualModelOverlay}: when the
* compaction-model recipe carries patch fields, synthesizes the derived
* registry entry ({@link COMPACTION_DERIVED_MODEL_ID}) into the effective
* `models` view — a copy of the pointed entry with the patch merged into its
* `overrides` block (patch wins conflicts) and `aliases` dropped, so the
* derived entry never competes in name/alias routing. Compaction-model binding
* then resolves it by name through the standard catalog path, and the patch
* rides the same `effectiveModelConfig` merge as any `models.*.overrides`
* (including its `supportEfforts` / `defaultEffort` pruning and input
* clamping).
*
* The synthesized entry lives ONLY in the in-memory effective view: `strip`
* removes it from `models` writes so it never reaches `config.toml`, and the
* persistence bridge's deep-equal guards keep the two-way sync silent. `strip`
* also rolls back a `defaultModel` pointer set to the derived id (restoring the
* raw value). Nothing is synthesized when the recipe has no patch fields (when
* `compaction.model` is unset), or when the pointed entry does not exist. The
* id is reserved: a user-configured entry under it is stripped on write all the
* same.
*
* Self-registered at module load via `registerConfigOverlay`; it is imported
* for side effects after the visual-model overlay.
*/
export const COMPACTION_DERIVED_MODEL_ID = '__compaction__';

export function compactionModelPatch(
compaction: CompactionModelConfig | undefined,
): ModelOverride | undefined {
if (compaction === undefined) return undefined;
const { model: _model, ...patch } = compaction;
return Object.keys(patch).length > 0 ? patch : undefined;
}

function asRecord(value: unknown): Record<string, unknown> {
return isPlainObject(value) ? value : {};
}

function withoutKey(value: unknown, key: string): unknown {
if (!isPlainObject(value) || !(key in value)) return value;
const out: Record<string, unknown> = { ...value };
delete out[key];
return out;
}

export const compactionModelOverlay: ConfigEffectiveOverlay = {
apply(effective, _getEnv, validate) {
const compaction = effective[COMPACTION_MODEL_SECTION] as
| CompactionModelConfig
| undefined;
const patch = compactionModelPatch(compaction);
const baseId = compaction?.model;
if (
patch === undefined ||
baseId === undefined ||
baseId === COMPACTION_DERIVED_MODEL_ID
) {
return [];
}
const models = asRecord(effective[MODELS_SECTION]);
const base = models[baseId];
if (!isPlainObject(base)) return [];
const { overrides: baseOverrides, aliases: _aliases, ...baseFields } = base;
const derived: Record<string, unknown> = {
...baseFields,
overrides: { ...asRecord(baseOverrides), ...patch },
};
effective[MODELS_SECTION] = validate(MODELS_SECTION, {
...models,
[COMPACTION_DERIVED_MODEL_ID]: derived,
});
return [MODELS_SECTION];
},

strip(domain, value, rawSnake) {
switch (domain) {
case MODELS_SECTION:
return withoutKey(value, COMPACTION_DERIVED_MODEL_ID);
case DEFAULT_MODEL_SECTION:
if (value !== COMPACTION_DERIVED_MODEL_ID) return value;
return typeof rawSnake['default_model'] === 'string'
? rawSnake['default_model']
: undefined;
default:
return value;
}
},
};

registerConfigOverlay(compactionModelOverlay);
22 changes: 22 additions & 0 deletions packages/agent-core-v2/src/app/kosongConfig/configSection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -335,6 +335,28 @@ registerConfigSection(VISUAL_MODEL_SECTION, VisualModelConfigSchema, {
});


export const COMPACTION_MODEL_SECTION = 'compactionModel';

export const COMPACTION_MODEL_ENV = 'KIMI_COMPACTION_MODEL';
export const COMPACTION_MODEL_EFFORT_ENV = 'KIMI_COMPACTION_EFFORT';

export const CompactionModelConfigSchema = ModelOverrideSchema.extend({
model: z.string().min(1).optional(),
});

export type CompactionModelConfig = z.infer<typeof CompactionModelConfigSchema>;

export const compactionModelEnvBindings = envBindings(CompactionModelConfigSchema, {
model: { env: COMPACTION_MODEL_ENV, parse: parseNonEmptyEnv },
defaultEffort: { env: COMPACTION_MODEL_EFFORT_ENV, parse: parseNonEmptyEnv },
});

registerConfigSection(COMPACTION_MODEL_SECTION, CompactionModelConfigSchema, {
env: compactionModelEnvBindings,
stripEnv: stripEnvBoundFields(compactionModelEnvBindings),
});


export const MODEL_CATALOG_SECTION = 'modelCatalog';

export const ModelCatalogConfigSchema = z.object({
Expand Down
4 changes: 4 additions & 0 deletions packages/agent-core-v2/src/app/telemetry/events.ts
Original file line number Diff line number Diff line change
Expand Up @@ -208,6 +208,8 @@ export interface CompactionFinishedEvent {
input_cache_read?: number;
input_cache_creation?: number;
trace_id?: string;
model?: string;
model_display?: string;
}

export interface CompactionFailedEvent {
Expand Down Expand Up @@ -676,6 +678,8 @@ export const telemetryEventDefinitions = {
input_cache_creation: 'Cache-creation input tokens',
trace_id:
'Trace id of the final compaction request round; absent for non-Kimi protocols',
model: 'Model alias that produced the compaction summary (dedicated compaction model when configured, otherwise the active conversation model)',
model_display: 'User-facing model alias for the compaction summary producer',
},
}),
compaction_failed: defineAgentTelemetryEvent<CompactionFailedEvent>({
Expand Down
25 changes: 24 additions & 1 deletion packages/agent-core-v2/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,7 @@ export * from '#/kosong/protocol/protocolBase';
export * from '#/kosong/protocol/protocolTrait';
import '#/app/kosongConfig/envOverlay';
import '#/app/kosongConfig/visualModelOverlay';
import '#/app/kosongConfig/compactionModelOverlay';
export * from '#/kosong/model/completionBudget';
export * from '#/kosong/model/hostRequestHeaders';
export * from '#/kosong/model/model';
Expand All @@ -190,7 +191,7 @@ export {
ModelCatalogConfigSchema,
type ModelCatalogConfig,
} from '#/app/kosongConfig/configSection';
export type { SecondaryModelConfig, VisualModelConfig } from '#/app/kosongConfig/configSection';
export type { SecondaryModelConfig, VisualModelConfig, CompactionModelConfig } from '#/app/kosongConfig/configSection';
export {
SECONDARY_MODEL_SECTION,
SECONDARY_MODEL_ENV,
Expand All @@ -202,12 +203,22 @@ export {
VISUAL_MODEL_EFFORT_ENV,
VisualModelConfigSchema,
visualModelEnvBindings,
COMPACTION_MODEL_SECTION,
COMPACTION_MODEL_ENV,
COMPACTION_MODEL_EFFORT_ENV,
CompactionModelConfigSchema,
compactionModelEnvBindings,
} from '#/app/kosongConfig/configSection';
export {
VISUAL_DERIVED_MODEL_ID,
visualModelOverlay,
visualModelPatch,
} from '#/app/kosongConfig/visualModelOverlay';
export {
COMPACTION_DERIVED_MODEL_ID,
compactionModelOverlay,
compactionModelPatch,
} from '#/app/kosongConfig/compactionModelOverlay';
export * from '#/app/kosongConfig/kosongConfig';
export * from '#/app/kosongConfig/kosongConfigService';
export * from '#/kosong/model/modelOAuth';
Expand Down Expand Up @@ -483,6 +494,8 @@ export * from '#/session/subagent/mirrorAgentRun';
import '#/session/subagent/configSection';
import '#/session/visual/flag';
import '#/session/visual/configSection';
import '#/session/compaction/flag';
import '#/session/compaction/configSection';
import '#/session/substitute/flag';
import '#/session/substitute/configSection';
export {
Expand All @@ -500,6 +513,16 @@ export {
VISUAL_MODEL_CHOICE_SCHEMA,
type VisualModelChoice,
} from '#/session/visual/configSection';
export {
COMPACTION_MODEL_FLAG_ID,
COMPACTION_MODEL_FLAG_ENV,
compactionModelFlag,
} from '#/session/compaction/flag';
export {
resolveCompactionModel,
resolveCompactionBinding,
compactionModelBindingFor,
} from '#/session/compaction/configSection';
export {
SUBSTITUTE_MODEL_FLAG_ID,
SUBSTITUTE_MODEL_FLAG_ENV,
Expand Down
Loading
Loading