fix(agent-core-v2): validate goal records (#1694)

Author: chengluyu

.changeset/validate-goal-records.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Reject malformed persisted goal records during session recovery.

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<GoalModelState>('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<GoalStatus>().optional(),
-    reason: z.string().optional(),
-    turnsUsed: z.number().optional(),
-    tokensUsed: z.number().optional(),
-    wallClockMs: z.number().optional(),
-    wallClockResumedAt: z.number().optional(),
-    budgetLimits: z.custom<GoalBudgetLimits>().optional(),
-    actor: z.custom<GoalActor>().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;

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

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);

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();
+    }
+  });
 });