From 2c908b497a6cf0cbe2b2d7d99c24124c2572e59e Mon Sep 17 00:00:00 2001 From: chengluyu <2239547+chengluyu@users.noreply.github.com> Date: Tue, 14 Jul 2026 17:14:23 +0800 Subject: [PATCH] fix(agent-core-v2): validate goal records --- .changeset/validate-goal-records.md | 5 + .../agent-core-v2/src/agent/goal/goalOps.ts | 53 +++++-- packages/agent-core-v2/src/wire/record.ts | 9 ++ .../agent-core-v2/src/wire/wireService.ts | 38 +++-- .../test/agent/goal/goalOps.test.ts | 149 ++++++++++++++++++ 5 files changed, 228 insertions(+), 26 deletions(-) create mode 100644 .changeset/validate-goal-records.md diff --git a/.changeset/validate-goal-records.md b/.changeset/validate-goal-records.md new file mode 100644 index 0000000000..ae07c164d1 --- /dev/null +++ b/.changeset/validate-goal-records.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Reject malformed persisted goal records during session recovery. diff --git a/packages/agent-core-v2/src/agent/goal/goalOps.ts b/packages/agent-core-v2/src/agent/goal/goalOps.ts index 70eaba0d29..953488d6de 100644 --- a/packages/agent-core-v2/src/agent/goal/goalOps.ts +++ b/packages/agent-core-v2/src/agent/goal/goalOps.ts @@ -32,7 +32,6 @@ import { z } from 'zod'; import { defineModel } from '#/wire/model'; import type { - GoalActor, GoalBudgetLimits, GoalChange, GoalSnapshot, @@ -56,6 +55,18 @@ export type GoalModelState = GoalState | null; export const GoalModel = defineModel('goal', () => null); +const GoalStatusSchema = z.enum(['active', 'paused', 'blocked', 'complete']); + +const GoalActorSchema = z.enum(['user', 'model', 'runtime', 'system']); + +const GoalBudgetLimitsSchema = z + .object({ + tokenBudget: z.number().finite().nonnegative().optional(), + turnBudget: z.number().finite().nonnegative().optional(), + wallClockBudgetMs: z.number().finite().nonnegative().optional(), + }) + .strict(); + declare module '#/app/event/eventBus' { interface DomainEventMap { 'goal.updated': { @@ -75,12 +86,17 @@ declare module '#/wire/types' { } export const createGoal = GoalModel.defineOp('goal.create', { - schema: z.object({ - goalId: z.string(), - objective: z.string(), - completionCriterion: z.string().optional(), - wallClockResumedAt: z.number().optional(), - }), + schema: z + .object({ + goalId: z.string(), + objective: z.string(), + completionCriterion: z.string().optional(), + wallClockResumedAt: z.number().finite().nonnegative().optional(), + status: GoalStatusSchema.optional(), + actor: GoalActorSchema.optional(), + budgetLimits: GoalBudgetLimitsSchema.optional(), + }) + .strip(), apply: (_s, p) => ({ goalId: p.goalId, objective: p.objective, @@ -95,16 +111,19 @@ export const createGoal = GoalModel.defineOp('goal.create', { }); export const updateGoal = GoalModel.defineOp('goal.update', { - schema: z.object({ - status: z.custom().optional(), - reason: z.string().optional(), - turnsUsed: z.number().optional(), - tokensUsed: z.number().optional(), - wallClockMs: z.number().optional(), - wallClockResumedAt: z.number().optional(), - budgetLimits: z.custom().optional(), - actor: z.custom().optional(), - }), + schema: z + .object({ + goalId: z.string().optional(), + status: GoalStatusSchema.optional(), + reason: z.string().optional(), + turnsUsed: z.number().finite().nonnegative().optional(), + tokensUsed: z.number().finite().nonnegative().optional(), + wallClockMs: z.number().finite().nonnegative().optional(), + wallClockResumedAt: z.number().finite().nonnegative().optional(), + budgetLimits: GoalBudgetLimitsSchema.optional(), + actor: GoalActorSchema.optional(), + }) + .strip(), apply: (s, p) => { if (s === null) return null; let next: GoalState | undefined; diff --git a/packages/agent-core-v2/src/wire/record.ts b/packages/agent-core-v2/src/wire/record.ts index 3a02da4eae..92a9bc28f9 100644 --- a/packages/agent-core-v2/src/wire/record.ts +++ b/packages/agent-core-v2/src/wire/record.ts @@ -25,6 +25,15 @@ export interface WireMetadataRecord extends WireRecord { readonly created_at: number; } +export function isWireRecord(record: unknown): record is WireRecord { + return ( + record !== null && + typeof record === 'object' && + !Array.isArray(record) && + typeof (record as { type?: unknown }).type === 'string' + ); +} + export function createWireMetadataRecord(now = Date.now()): WireMetadataRecord { return { type: 'metadata', diff --git a/packages/agent-core-v2/src/wire/wireService.ts b/packages/agent-core-v2/src/wire/wireService.ts index 18cc5ec9ce..a02a8441b6 100644 --- a/packages/agent-core-v2/src/wire/wireService.ts +++ b/packages/agent-core-v2/src/wire/wireService.ts @@ -41,6 +41,7 @@ import { OP_REGISTRY } from './op'; import { AGENT_WIRE_RECORD_KEY, createWireMetadataRecord, + isWireRecord, isWireMetadataRecord, opToWireRecord, wireRecordToPayload, @@ -149,7 +150,13 @@ export class WireService extends Disposable implements IWireService { let recordIndex = 0; let hasRecords = false; - for await (const sourceRecord of source) { + for await (const candidate of source) { + const sourceRecord: unknown = candidate; + if (!isWireRecord(sourceRecord)) { + this.reportSkippedRecord(undefined, recordIndex, true); + recordIndex++; + continue; + } if (!hasRecords) { hasRecords = true; if (sourceRecord.type !== 'metadata') { @@ -207,21 +214,34 @@ export class WireService extends Disposable implements IWireService { private replayRecord(record: WireRecord, index: number): void { const descriptor = OP_REGISTRY.get(record.type); if (descriptor === undefined) { - onUnexpectedError( - new WireError( - WireErrors.codes.WIRE_UNKNOWN_RECORD, - `Unknown wire record type '${record.type}' skipped during restore`, - { details: { type: record.type, index } }, - ), - ); + this.reportSkippedRecord(record.type, index); + return; + } + const payload = descriptor.schema.safeParse(wireRecordToPayload(record)); + if (!payload.success) { + this.reportSkippedRecord(record.type, index, true); return; } this.execute({ - ops: [{ type: record.type, payload: wireRecordToPayload(record), descriptor }], + ops: [{ type: record.type, payload: payload.data, descriptor }], silent: true, }); } + private reportSkippedRecord(type: string | undefined, index: number, malformed = false): void { + onUnexpectedError( + new WireError( + WireErrors.codes.WIRE_UNKNOWN_RECORD, + type === undefined + ? 'Malformed wire record skipped during restore' + : malformed + ? `Malformed wire record type '${type}' skipped during restore` + : `Unknown wire record type '${type}' skipped during restore`, + { details: { type, index } }, + ), + ); + } + private execute(group: OpGroup): void { for (const op of group.ops) { const inst = this.ensureModel(op.descriptor.model); diff --git a/packages/agent-core-v2/test/agent/goal/goalOps.test.ts b/packages/agent-core-v2/test/agent/goal/goalOps.test.ts index 5746cc3373..f5f6192ea9 100644 --- a/packages/agent-core-v2/test/agent/goal/goalOps.test.ts +++ b/packages/agent-core-v2/test/agent/goal/goalOps.test.ts @@ -9,6 +9,7 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { SyncDescriptor } from '#/_base/di/descriptors'; import { DisposableStore } from '#/_base/di/lifecycle'; import { TestInstantiationService } from '#/_base/di/test'; +import { resetUnexpectedErrorHandler, setUnexpectedErrorHandler } from '#/_base/errors/unexpectedError'; import { Event } from '#/_base/event'; import { IEventBus } from '#/app/event/eventBus'; import { EventBusService } from '#/app/event/eventBusService'; @@ -264,4 +265,152 @@ describe('AgentGoalService (wire-backed)', () => { }), ]); }); + + it('restores goal records with omitted optional fields from older journals', async () => { + await restoreTestAgentWire(wire, log, testWireScope(SCOPE, KEY), [ + { type: 'goal.create', goalId: 'goal-1', objective: 'work' }, + { type: 'goal.update' }, + ]); + + expect(modelOf(wire)).toMatchObject({ + goalId: 'goal-1', + status: 'paused', + budgetLimits: {}, + }); + }); + + it('restores legacy goal create audit fields without changing normalized state', async () => { + await restoreTestAgentWire(wire, log, testWireScope(SCOPE, KEY), [ + { + type: 'goal.create', + goalId: 'goal-1', + objective: 'work', + status: 'active', + actor: 'user', + budgetLimits: {}, + }, + ]); + + expect(modelOf(wire)).toMatchObject({ + goalId: 'goal-1', + status: 'paused', + budgetLimits: {}, + }); + }); + + it('restores a legacy goal update identity without changing state selection', async () => { + await restoreTestAgentWire(wire, log, testWireScope(SCOPE, KEY), [ + { type: 'goal.create', goalId: 'goal-1', objective: 'work' }, + { type: 'goal.update', goalId: 'goal-1', status: 'blocked', reason: 'waiting' }, + ]); + + expect(modelOf(wire)).toMatchObject({ + goalId: 'goal-1', + status: 'blocked', + terminalReason: 'waiting', + }); + }); + + it('strips forward-compatible goal fields during restore', async () => { + await restoreTestAgentWire(wire, log, testWireScope(SCOPE, KEY), [ + { + type: 'goal.create', + goalId: 'goal-1', + objective: 'work', + futureField: true, + }, + ]); + + expect(modelOf(wire)).toMatchObject({ goalId: 'goal-1', objective: 'work' }); + }); + + it('skips a goal update with an invalid status during restore', async () => { + const unexpected: unknown[] = []; + setUnexpectedErrorHandler((error) => unexpected.push(error)); + try { + await restoreTestAgentWire(wire, log, testWireScope(SCOPE, KEY), [ + { type: 'goal.create', goalId: 'goal-1', objective: 'work' }, + { type: 'goal.update', status: 'cancelled' }, + ]); + + expect(modelOf(wire)).toMatchObject({ status: 'paused' }); + expect(unexpected).toContainEqual( + expect.objectContaining({ code: 'wire.unknown_record', details: { type: 'goal.update', index: 1 } }), + ); + } finally { + resetUnexpectedErrorHandler(); + } + }); + + it('skips a goal update with an invalid actor during restore', async () => { + const unexpected: unknown[] = []; + setUnexpectedErrorHandler((error) => unexpected.push(error)); + try { + await restoreTestAgentWire(wire, log, testWireScope(SCOPE, KEY), [ + { type: 'goal.create', goalId: 'goal-1', objective: 'work' }, + { type: 'goal.update', actor: 'assistant' }, + ]); + + expect(modelOf(wire)).toMatchObject({ status: 'paused' }); + expect(unexpected).toContainEqual( + expect.objectContaining({ code: 'wire.unknown_record', details: { type: 'goal.update', index: 1 } }), + ); + } finally { + resetUnexpectedErrorHandler(); + } + }); + + it('skips negative and non-finite goal counters and budgets during restore', async () => { + const unexpected: unknown[] = []; + setUnexpectedErrorHandler((error) => unexpected.push(error)); + try { + await restoreTestAgentWire(wire, log, testWireScope(SCOPE, KEY), [ + { type: 'goal.create', goalId: 'goal-1', objective: 'work' }, + { type: 'goal.update', turnsUsed: -1 }, + { type: 'goal.update', tokensUsed: Number.POSITIVE_INFINITY }, + { type: 'goal.update', wallClockMs: Number.NaN }, + { type: 'goal.update', wallClockResumedAt: Number.NaN }, + { type: 'goal.update', budgetLimits: { turnBudget: -1 } }, + { type: 'goal.update', budgetLimits: { tokenBudget: Number.POSITIVE_INFINITY } }, + { type: 'goal.update', budgetLimits: { wallClockBudgetMs: Number.NaN } }, + ]); + + expect(modelOf(wire)).toMatchObject({ + turnsUsed: 0, + tokensUsed: 0, + wallClockMs: 0, + budgetLimits: {}, + }); + expect(unexpected).toHaveLength(7); + } finally { + resetUnexpectedErrorHandler(); + } + }); + + it('skips null, arrays, and malformed nested goal records during restore', async () => { + const unexpected: unknown[] = []; + setUnexpectedErrorHandler((error) => unexpected.push(error)); + try { + await restoreTestAgentWire( + wire, + log, + testWireScope(SCOPE, KEY), + [ + null, + [], + { + type: 'goal.create', + goalId: 'goal-1', + objective: 'work', + budgetLimits: { unexpected: true }, + }, + ] as unknown as WireRecord[], + ); + + expect(modelOf(wire)).toBeNull(); + expect(unexpected).toHaveLength(3); + } finally { + resetUnexpectedErrorHandler(); + } + }); });