refactor(v2): bind op definitions to models and require zod payload schemas

.changeset/v2-wire-op-schemas.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Declare v2 engine wire op payloads with required zod schemas and derive their types from the schemas, with ops declared on their models and every op type registered for replay classification.

packages/agent-core-v2/src/activity/activityOps.ts

@@ -14,8 +14,9 @@
  * Consumed by the Agent-scope `agentActivityService` and (PR5) the projector.
  */
 
+import { z } from 'zod';
+
 import { defineModel } from '#/wire/model';
-import { defineOp } from '#/wire/op';
 import type { PromptOrigin } from '#/agent/contextMemory/types';
 
 import type { AgentLane, BackgroundActivityRef, SessionLane } from './activity';
@@ -46,10 +47,17 @@ export const LaneModel = defineModel<LaneModelState>('activityLane', () => ({
   background: [],
 }));
 
-export const setLane = defineOp(LaneModel, 'activity.set_lane', {
+declare module '#/wire/types' {
+  interface TransientOpMap {
+    'activity.set_lane': typeof setLane;
+    'activity.set_session_lane': typeof setSessionLane;
+  }
+}
+
+export const setLane = LaneModel.defineOp('activity.set_lane', {
+  schema: z.object({ next: z.custom<LaneModelState>() }),
   persist: false,
-  apply: (s, p: { next: LaneModelState }): LaneModelState =>
-    laneEqual(s, p.next) ? s : p.next,
+  apply: (s, p) => (laneEqual(s, p.next) ? s : p.next),
 });
 
 export function laneEqual(a: LaneModelState, b: LaneModelState): boolean {
@@ -84,8 +92,8 @@ export const SessionLaneModel = defineModel<SessionLaneModelState>('sessionActiv
   activeLeases: 0,
 }));
 
-export const setSessionLane = defineOp(SessionLaneModel, 'activity.set_session_lane', {
+export const setSessionLane = SessionLaneModel.defineOp('activity.set_session_lane', {
+  schema: z.object({ next: z.custom<SessionLaneModelState>() }),
   persist: false,
-  apply: (s, p: { next: SessionLaneModelState }): SessionLaneModelState =>
-    s.lane === p.next.lane && s.activeLeases === p.next.activeLeases ? s : p.next,
+  apply: (s, p) => (s.lane === p.next.lane && s.activeLeases === p.next.activeLeases ? s : p.next),
 });

packages/agent-core-v2/src/agent/contextMemory/contextOps.ts

@@ -32,9 +32,10 @@
  *   data that was compacted away during the session.
  */
 
+import { z } from 'zod';
+
 import type { ContentPart } from '#/app/llmProtocol/message';
 import { defineModel, type PartsTransformer } from '#/wire/model';
-import { defineOp } from '#/wire/op';
 import type { PersistedRecord } from '#/wire/wireService';
 
 import {
@@ -120,63 +121,71 @@ function popSwarmModeReminder(state: ContextMessage[], _payload: unknown): Conte
   return resetFold(state.slice(0, -1)) as ContextMessage[];
 }
 
-export interface ContextMessagePayload {
-  readonly message: ContextMessage;
+declare module '#/wire/types' {
+  interface PersistedOpMap {
+    'context.append_message': typeof contextAppendMessage;
+    'context.append_loop_event': typeof contextAppendLoopEvent;
+    'context.clear': typeof contextClear;
+    'context.apply_compaction': typeof contextApplyCompaction;
+    'context.undo': typeof contextUndo;
+  }
 }
 
-export const contextAppendMessage = defineOp(ContextModel, 'context.append_message', {
-  apply: (state, p: ContextMessagePayload): ContextMessage[] =>
-    foldAppendMessage(state, p.message) as ContextMessage[],
-});
-
-export interface ContextLoopEventPayload {
-  readonly event: LoopRecordedEvent;
-}
+// `ContextMessage` / `LoopRecordedEvent` are large domain unions owned by
+// sibling modules; `z.custom` keeps their exact types without restating them.
+const contextMessageSchema = z.custom<ContextMessage>();
+const loopRecordedEventSchema = z.custom<LoopRecordedEvent>();
 
-export const contextAppendLoopEvent = defineOp(ContextModel, 'context.append_loop_event', {
-  apply: (state, p: ContextLoopEventPayload): ContextMessage[] =>
-    foldLoopEvent(state, p.event) as ContextMessage[],
+export const contextAppendMessage = ContextModel.defineOp('context.append_message', {
+  schema: z.object({ message: contextMessageSchema }),
+  apply: (state, p) => foldAppendMessage(state, p.message) as ContextMessage[],
 });
 
-export const contextClear = defineOp(ContextModel, 'context.clear', {
-  apply: (state): ContextMessage[] =>
-    state.length === 0 ? state : (resetFold([]) as ContextMessage[]),
+export const contextAppendLoopEvent = ContextModel.defineOp('context.append_loop_event', {
+  schema: z.object({ event: loopRecordedEventSchema }),
+  apply: (state, p) => foldLoopEvent(state, p.event) as ContextMessage[],
 });
 
-interface ContextCompactionBasePayload {
-  readonly tokensBefore?: number;
-  readonly tokensAfter?: number;
-  readonly keptUserMessageCount?: number;
-  readonly keptHeadUserMessageCount?: number;
-  readonly droppedCount?: number;
-  readonly legacyTail?: boolean;
-}
-
-export interface TextSummaryCompactionPayload extends ContextCompactionBasePayload {
-  readonly summary: string;
-  readonly compactedCount: number;
-  readonly contextSummary?: string;
-}
-
-interface ContextSummaryCompactionPayload extends ContextCompactionBasePayload {
-  readonly contextSummary: string;
-  readonly compactedCount: number;
-  readonly summary?: string;
-}
-
-interface LegacyMessageSummaryCompactionPayload extends ContextCompactionBasePayload {
-  readonly summary: ContextMessage;
-  readonly count: number;
-  readonly compactedCount?: number;
-}
-
-export type ContextCompactionPayload =
-  | TextSummaryCompactionPayload
-  | ContextSummaryCompactionPayload
-  | LegacyMessageSummaryCompactionPayload;
+export const contextClear = ContextModel.defineOp('context.clear', {
+  schema: z.object({}),
+  apply: (state) => (state.length === 0 ? state : (resetFold([]) as ContextMessage[])),
+});
 
-export const contextApplyCompaction = defineOp(ContextModel, 'context.apply_compaction', {
-  apply: (state, p: ContextCompactionPayload): ContextMessage[] => {
+const contextCompactionBaseShape = {
+  tokensBefore: z.number().optional(),
+  tokensAfter: z.number().optional(),
+  keptUserMessageCount: z.number().optional(),
+  keptHeadUserMessageCount: z.number().optional(),
+  droppedCount: z.number().optional(),
+  legacyTail: z.boolean().optional(),
+};
+
+const contextApplyCompactionSchema = z.union([
+  z.object({
+    ...contextCompactionBaseShape,
+    summary: z.string(),
+    compactedCount: z.number(),
+    contextSummary: z.string().optional(),
+  }),
+  z.object({
+    ...contextCompactionBaseShape,
+    contextSummary: z.string(),
+    compactedCount: z.number(),
+    summary: z.string().optional(),
+  }),
+  z.object({
+    ...contextCompactionBaseShape,
+    summary: contextMessageSchema,
+    count: z.number(),
+    compactedCount: z.number().optional(),
+  }),
+]);
+
+type ContextCompactionPayload = z.infer<typeof contextApplyCompactionSchema>;
+
+export const contextApplyCompaction = ContextModel.defineOp('context.apply_compaction', {
+  schema: contextApplyCompactionSchema,
+  apply: (state, p) => {
     const result = buildContextCompactionShape(state, readContextCompactionShapeInput(p));
     return resetFold([...result.messages]) as ContextMessage[];
   },
@@ -279,10 +288,6 @@ function isContextMessage(value: unknown): value is ContextMessage {
   return typeof message.role === 'string' && Array.isArray(message.content);
 }
 
-export interface ContextUndoPayload {
-  readonly count: number;
-}
-
 export interface UndoCut {
   readonly cutIndex: number;
   readonly removedCount: number;
@@ -370,8 +375,9 @@ export function formatUndoUnavailableMessage(
   }
 }
 
-export const contextUndo = defineOp(ContextModel, 'context.undo', {
-  apply: (state, p: ContextUndoPayload): ContextMessage[] => {
+export const contextUndo = ContextModel.defineOp('context.undo', {
+  schema: z.object({ count: z.number() }),
+  apply: (state, p) => {
     if (p.count <= 0 || state.length === 0) return state;
     const cut = computeUndoCut(state, p.count);
     if (!isFullyUndoable(cut, p.count)) return state;

packages/agent-core-v2/src/agent/contextSize/contextSizeOps.ts

@@ -22,8 +22,9 @@
  * `contextSizeService`.
  */
 
+import { z } from 'zod';
+
 import { defineModel } from '#/wire/model';
-import { defineOp } from '#/wire/op';
 
 export interface ContextSizeState {
   readonly length: number;
@@ -35,14 +36,16 @@ export const ContextSizeModel = defineModel<ContextSizeState>('contextSize', ()
   tokens: 0,
 }));
 
-export interface ContextSizeMeasuredPayload {
-  readonly length: number;
-  readonly tokens: number;
+declare module '#/wire/types' {
+  interface TransientOpMap {
+    'context_size.measured': typeof contextSizeMeasured;
+  }
 }
 
-export const contextSizeMeasured = defineOp(ContextSizeModel, 'context_size.measured', {
+export const contextSizeMeasured = ContextSizeModel.defineOp('context_size.measured', {
+  schema: z.object({ length: z.number(), tokens: z.number() }),
   persist: false,
-  apply: (s, p: ContextSizeMeasuredPayload): ContextSizeState => {
+  apply: (s, p) => {
     const length = normalizeMeasuredLength(p.length);
     const tokens = Math.max(0, p.tokens);
     if (s.length === length && s.tokens === tokens) return s;

packages/agent-core-v2/src/agent/contextSize/contextSizeService.ts

@@ -4,9 +4,9 @@
  * Owns the last measured context token count in the wire `ContextSizeModel`
  * (`{ length, tokens }`): reads it through `wire.getModel`, writes it through
  * `wire.dispatch(contextSizeMeasured(...))` (called by `llmRequester` after each
- * measured exchange), and emits the `contextTokens` slice of
- * `agent.status.updated` live through `wire.signal` when the measured value
- * changes. `get(start?, end?)` returns `{ size, measured, estimated }` for the
+ * measured exchange), and derives the `contextTokens` slice of
+ * `agent.status.updated` from the Op's `toEvent` (published to `IEventBus` on
+ * dispatch) when the measured value changes. `get(start?, end?)` returns `{ size, measured, estimated }` for the
  * context-message range `[start, end)`, resolved like `Array.prototype.slice`
  * (defaulting to the whole context; negative indices count back from the end;
  * an inverted range is empty): `measured`

packages/agent-core-v2/src/agent/fullCompaction/compactionOps.ts

@@ -27,17 +27,17 @@
  * `wire.onRestored` handler (mirroring `goal`'s post-replay normalization).
  *
  * The `compaction.*` events publish to `IEventBus` (`compaction.started` via the
- * `begin` Op's `toEvent`; the rest directly) and also emit live through
- * `wire.signal` (legacy channel, until Phase 3); they are declared here via
- * interface-merge (`error` is already declared by `mcp`, so it is not
- * re-declared). The `full_compaction.*` record shapes stay declared in
- * `WireRecordMap` (see `fullCompactionService.ts`) because the records still
+ * `begin` Op's `toEvent`; the rest directly from the service); they are
+ * declared here via interface-merge (`error` is already declared by `mcp`, so
+ * it is not re-declared). The `full_compaction.*` record shapes are registered in
+ * `PersistedOpMap` (`#/wire/types`, below) because the records still
  * ride the per-agent `wire.jsonl` log read by `wireRecord.restore()` /
  * `getRecords()`. Consumed by the Agent-scope `fullCompactionService`.
  */
 
+import { z } from 'zod';
+
 import { defineModel } from '#/wire/model';
-import { defineOp } from '#/wire/op';
 import type {
   CompactionBlockedEvent,
   CompactionCancelledEvent,
@@ -66,25 +66,30 @@ declare module '#/app/event/eventBus' {
   }
 }
 
-export type FullCompactionBeginPayload = CompactionBeginData;
+declare module '#/wire/types' {
+  interface PersistedOpMap {
+    'full_compaction.begin': typeof fullCompactionBegin;
+    'full_compaction.cancel': typeof fullCompactionCancel;
+    'full_compaction.complete': typeof fullCompactionComplete;
+  }
+}
 
-export const fullCompactionBegin = defineOp(CompactionModel, 'full_compaction.begin', {
-  apply: (s, _p: FullCompactionBeginPayload): CompactionState =>
-    s.phase === 'running' ? s : { phase: 'running' },
+export const fullCompactionBegin = CompactionModel.defineOp('full_compaction.begin', {
+  schema: z.custom<CompactionBeginData>(),
+  apply: (s) => (s.phase === 'running' ? s : { phase: 'running' }),
   toEvent: (p) => ({
     type: 'compaction.started' as const,
     trigger: p.source,
     instruction: p.instruction,
   }),
 });
 
-export const fullCompactionCancel = defineOp(CompactionModel, 'full_compaction.cancel', {
-  apply: (s): CompactionState => (s.phase === 'idle' ? s : { phase: 'idle' }),
+export const fullCompactionCancel = CompactionModel.defineOp('full_compaction.cancel', {
+  schema: z.object({}),
+  apply: (s) => (s.phase === 'idle' ? s : { phase: 'idle' }),
 });
 
-export type FullCompactionCompletePayload = Record<string, never>;
-
-export const fullCompactionComplete = defineOp(CompactionModel, 'full_compaction.complete', {
-  apply: (s, _p: FullCompactionCompletePayload): CompactionState =>
-    s.phase === 'idle' ? s : { phase: 'idle' },
+export const fullCompactionComplete = CompactionModel.defineOp('full_compaction.complete', {
+  schema: z.object({}),
+  apply: (s) => (s.phase === 'idle' ? s : { phase: 'idle' }),
 });

packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts

@@ -56,7 +56,6 @@ import {
   fullCompactionBegin,
   fullCompactionCancel,
   fullCompactionComplete,
-  type FullCompactionCompletePayload,
 } from './compactionOps';
 import {
   type CompactionBeginData,
@@ -65,19 +64,6 @@ import {
 import { Emitter, type Event } from '#/_base/event';
 import { OrderedHookSlot } from '#/hooks';
 
-// The `full_compaction.*` record shapes stay declared in `WireRecordMap`
-// because the records still ride the per-agent `wire.jsonl` log read by
-// `wireRecord.restore()` / `getRecords()`. fullCompaction itself no longer
-// registers resumers here — its state rebuilds from the same log via
-// `wire.replay` into `CompactionModel`.
-declare module '#/agent/wireRecord/wireRecord' {
-  interface WireRecordMap {
-    'full_compaction.begin': CompactionBeginData;
-    'full_compaction.cancel': {};
-    'full_compaction.complete': FullCompactionCompletePayload;
-  }
-}
-
 export const MAX_COMPACTION_RETRY_ATTEMPTS = 5;
 const DEFAULT_COMPACTION_MAX_COMPLETION_TOKENS = 128 * 1024;
 const OVERFLOW_CONTEXT_SAFETY_RATIO = 0.85;