refactor(core): migrate MessageV2 message DTOs (User/Assistant/Part/Info/WithParts) to Effect Schema (#23757)

2026-04-30 13:51:48 +08:00 · 2026-04-21 23:26:12 -04:00
parent 1a76799fd8
commit e89543811c
6 changed files with 132 additions and 116 deletions
--- a/packages/opencode/src/cli/cmd/import.ts
+++ b/packages/opencode/src/cli/cmd/import.ts
@@ -168,7 +168,7 @@ export const ImportCommand = cmd({
      )

      for (const msg of exportData.messages) {
-        const msgInfo = MessageV2.Info.parse(msg.info)
+        const msgInfo = MessageV2.Info.zod.parse(msg.info)
        const { id, sessionID: _, ...msgData } = msgInfo
        Database.use((db) =>
          db
@@ -184,7 +184,7 @@ export const ImportCommand = cmd({
        )

        for (const part of msg.parts) {
-          const partInfo = MessageV2.Part.parse(part)
+          const partInfo = MessageV2.Part.zod.parse(part)
          const { id: partId, sessionID: _s, messageID, ...partData } = partInfo
          Database.use((db) =>
            db
--- a/packages/opencode/src/server/routes/instance/session.ts
+++ b/packages/opencode/src/server/routes/instance/session.ts
@@ -611,7 +611,7 @@ export const SessionRoutes = lazy(() =>
            description: "List of messages",
            content: {
              "application/json": {
-                schema: resolver(MessageV2.WithParts.array()),
+                schema: resolver(MessageV2.WithParts.zod.array()),
              },
            },
          },
@@ -701,8 +701,8 @@ export const SessionRoutes = lazy(() =>
              "application/json": {
                schema: resolver(
                  z.object({
-                    info: MessageV2.Info,
-                    parts: MessageV2.Part.array(),
+                    info: MessageV2.Info.zod,
+                    parts: MessageV2.Part.zod.array(),
                  }),
                ),
              },
@@ -813,7 +813,7 @@ export const SessionRoutes = lazy(() =>
            description: "Successfully updated part",
            content: {
              "application/json": {
-                schema: resolver(MessageV2.Part),
+                schema: resolver(MessageV2.Part.zod),
              },
            },
          },
@@ -828,7 +828,7 @@ export const SessionRoutes = lazy(() =>
          partID: PartID.zod,
        }),
      ),
-      validator("json", MessageV2.Part),
+      validator("json", MessageV2.Part.zod),
      async (c) => {
        const params = c.req.valid("param")
        const body = c.req.valid("json")
@@ -856,8 +856,8 @@ export const SessionRoutes = lazy(() =>
              "application/json": {
                schema: resolver(
                  z.object({
-                    info: MessageV2.Assistant,
-                    parts: MessageV2.Part.array(),
+                    info: MessageV2.Assistant.zod,
+                    parts: MessageV2.Part.zod.array(),
                  }),
                ),
              },
@@ -944,8 +944,8 @@ export const SessionRoutes = lazy(() =>
              "application/json": {
                schema: resolver(
                  z.object({
-                    info: MessageV2.Assistant,
-                    parts: MessageV2.Part.array(),
+                    info: MessageV2.Assistant.zod,
+                    parts: MessageV2.Part.zod.array(),
                  }),
                ),
              },
@@ -980,7 +980,7 @@ export const SessionRoutes = lazy(() =>
            description: "Created message",
            content: {
              "application/json": {
-                schema: resolver(MessageV2.WithParts),
+                schema: resolver(MessageV2.WithParts.zod),
              },
            },
          },
--- a/packages/opencode/src/session/message-v2.ts
+++ b/packages/opencode/src/session/message-v2.ts
@@ -368,37 +368,68 @@ export type ToolPart = Omit<Types.DeepMutable<Schema.Schema.Type<typeof ToolPart
  state: ToolState
 }

-const Base = z.object({
-  id: MessageID.zod,
-  sessionID: SessionID.zod,
-})
+const messageBase = {
+  id: MessageID,
+  sessionID: SessionID,
+}

-export const User = Base.extend({
-  role: z.literal("user"),
-  time: z.object({
-    created: z.number(),
+export const User = Schema.Struct({
+  ...messageBase,
+  role: Schema.Literal("user"),
+  time: Schema.Struct({
+    created: Schema.Number,
  }),
-  format: Format.zod.optional(),
-  summary: z
-    .object({
-      title: z.string().optional(),
-      body: z.string().optional(),
-      diffs: Snapshot.FileDiff.zod.array(),
-    })
-    .optional(),
-  agent: z.string(),
-  model: z.object({
-    providerID: ProviderID.zod,
-    modelID: ModelID.zod,
-    variant: z.string().optional(),
+  format: Schema.optional(_Format),
+  summary: Schema.optional(
+    Schema.Struct({
+      title: Schema.optional(Schema.String),
+      body: Schema.optional(Schema.String),
+      diffs: Schema.Array(Snapshot.FileDiff),
+    }),
+  ),
+  agent: Schema.String,
+  model: Schema.Struct({
+    providerID: ProviderID,
+    modelID: ModelID,
+    variant: Schema.optional(Schema.String),
  }),
-  system: z.string().optional(),
-  tools: z.record(z.string(), z.boolean()).optional(),
-}).meta({
-  ref: "UserMessage",
+  system: Schema.optional(Schema.String),
+  tools: Schema.optional(Schema.Record(Schema.String, Schema.Boolean)),
 })
-export type User = z.infer<typeof User>
+  .annotate({ identifier: "UserMessage" })
+  .pipe(withStatics((s) => ({ zod: zod(s) })))
+export type User = Types.DeepMutable<Schema.Schema.Type<typeof User>>

+const _Part = Schema.Union([
+  TextPart,
+  SubtaskPart,
+  ReasoningPart,
+  FilePart,
+  ToolPart,
+  StepStartPart,
+  StepFinishPart,
+  SnapshotPart,
+  PatchPart,
+  AgentPart,
+  RetryPart,
+  CompactionPart,
+]).annotate({ discriminator: "type", identifier: "Part" })
+export const Part = Object.assign(_Part, {
+  zod: zod(_Part) as unknown as z.ZodType<
+    | TextPart
+    | SubtaskPart
+    | ReasoningPart
+    | FilePart
+    | ToolPart
+    | StepStartPart
+    | StepFinishPart
+    | SnapshotPart
+    | PatchPart
+    | AgentPart
+    | RetryPart
+    | CompactionPart
+  >,
+})
 export type Part =
  | TextPart
  | SubtaskPart
@@ -413,28 +444,19 @@ export type Part =
  | RetryPart
  | CompactionPart

-// The derived `.zod` on each leaf is typed as `z.ZodType<...>`, but the walker
-// always emits a `z.ZodObject` at runtime. `z.discriminatedUnion` and
-// `z.infer` both rely on the ZodObject structural type, so cast here so the
-// resulting Part behaves like the pre-migration Zod union.
-export const Part = z
-  .discriminatedUnion("type", [
-    TextPart.zod as unknown as z.ZodObject<any>,
-    SubtaskPart.zod as unknown as z.ZodObject<any>,
-    ReasoningPart.zod as unknown as z.ZodObject<any>,
-    FilePart.zod as unknown as z.ZodObject<any>,
-    ToolPart.zod as unknown as z.ZodObject<any>,
-    StepStartPart.zod as unknown as z.ZodObject<any>,
-    StepFinishPart.zod as unknown as z.ZodObject<any>,
-    SnapshotPart.zod as unknown as z.ZodObject<any>,
-    PatchPart.zod as unknown as z.ZodObject<any>,
-    AgentPart.zod as unknown as z.ZodObject<any>,
-    RetryPart.zod as unknown as z.ZodObject<any>,
-    CompactionPart.zod as unknown as z.ZodObject<any>,
-  ])
-  .meta({
-    ref: "Part",
-  }) as unknown as z.ZodType<Part>
+// Errors are still NamedError-based Zod; bridge via ZodOverride so the derived
+// Zod + JSON Schema emit the original discriminatedUnion shape. Migrating the
+// error classes to Schema.TaggedErrorClass is a separate slice.
+const AssistantErrorZod = z.discriminatedUnion("name", [
+  AuthError.Schema,
+  NamedError.Unknown.Schema,
+  OutputLengthError.Schema,
+  AbortedError.Schema,
+  StructuredOutputError.Schema,
+  ContextOverflowError.Schema,
+  APIError.Schema,
+])
+type AssistantError = z.infer<typeof AssistantErrorZod>

 // ── Prompt input schemas ─────────────────────────────────────────────────────
 //
@@ -508,59 +530,53 @@ export const SubtaskPartInput = Schema.Struct({
  .pipe(withStatics((s) => ({ zod: zod(s) })))
 export type SubtaskPartInput = Types.DeepMutable<Schema.Schema.Type<typeof SubtaskPartInput>>

-export const Assistant = Base.extend({
-  role: z.literal("assistant"),
-  time: z.object({
-    created: z.number(),
-    completed: z.number().optional(),
+export const Assistant = Schema.Struct({
+  ...messageBase,
+  role: Schema.Literal("assistant"),
+  time: Schema.Struct({
+    created: Schema.Number,
+    completed: Schema.optional(Schema.Number),
  }),
-  error: z
-    .discriminatedUnion("name", [
-      AuthError.Schema,
-      NamedError.Unknown.Schema,
-      OutputLengthError.Schema,
-      AbortedError.Schema,
-      StructuredOutputError.Schema,
-      ContextOverflowError.Schema,
-      APIError.Schema,
-    ])
-    .optional(),
-  parentID: MessageID.zod,
-  modelID: ModelID.zod,
-  providerID: ProviderID.zod,
+  error: Schema.optional(Schema.Any.annotate({ [ZodOverride]: AssistantErrorZod })),
+  parentID: MessageID,
+  modelID: ModelID,
+  providerID: ProviderID,
  /**
   * @deprecated
   */
-  mode: z.string(),
-  agent: z.string(),
-  path: z.object({
-    cwd: z.string(),
-    root: z.string(),
+  mode: Schema.String,
+  agent: Schema.String,
+  path: Schema.Struct({
+    cwd: Schema.String,
+    root: Schema.String,
  }),
-  summary: z.boolean().optional(),
-  cost: z.number(),
-  tokens: z.object({
-    total: z.number().optional(),
-    input: z.number(),
-    output: z.number(),
-    reasoning: z.number(),
-    cache: z.object({
-      read: z.number(),
-      write: z.number(),
+  summary: Schema.optional(Schema.Boolean),
+  cost: Schema.Number,
+  tokens: Schema.Struct({
+    total: Schema.optional(Schema.Number),
+    input: Schema.Number,
+    output: Schema.Number,
+    reasoning: Schema.Number,
+    cache: Schema.Struct({
+      read: Schema.Number,
+      write: Schema.Number,
    }),
  }),
-  structured: z.any().optional(),
-  variant: z.string().optional(),
-  finish: z.string().optional(),
-}).meta({
-  ref: "AssistantMessage",
+  structured: Schema.optional(Schema.Any),
+  variant: Schema.optional(Schema.String),
+  finish: Schema.optional(Schema.String),
 })
-export type Assistant = z.infer<typeof Assistant>
+  .annotate({ identifier: "AssistantMessage" })
+  .pipe(withStatics((s) => ({ zod: zod(s) })))
+export type Assistant = Omit<Types.DeepMutable<Schema.Schema.Type<typeof Assistant>>, "error"> & {
+  error?: AssistantError
+}

-export const Info = z.discriminatedUnion("role", [User, Assistant]).meta({
-  ref: "Message",
+const _Info = Schema.Union([User, Assistant]).annotate({ discriminator: "role", identifier: "Message" })
+export const Info = Object.assign(_Info, {
+  zod: zod(_Info) as unknown as z.ZodType<User | Assistant>,
 })
-export type Info = z.infer<typeof Info>
+export type Info = User | Assistant

 export const Event = {
  Updated: SyncEvent.define({
@@ -569,7 +585,7 @@ export const Event = {
    aggregate: "sessionID",
    schema: z.object({
      sessionID: SessionID.zod,
-      info: Info,
+      info: Info.zod,
    }),
  }),
  Removed: SyncEvent.define({
@@ -587,7 +603,7 @@ export const Event = {
    aggregate: "sessionID",
    schema: z.object({
      sessionID: SessionID.zod,
-      part: Part,
+      part: Part.zod,
      time: z.number(),
    }),
  }),
@@ -613,10 +629,10 @@ export const Event = {
  }),
 }

-export const WithParts = z.object({
-  info: Info,
-  parts: z.array(Part),
-})
+export const WithParts = Schema.Struct({
+  info: _Info,
+  parts: Schema.Array(_Part),
+}).pipe(withStatics((s) => ({ zod: zod(s) })))
 export type WithParts = {
  info: Info
  parts: Part[]
--- a/packages/opencode/src/session/prompt.ts
+++ b/packages/opencode/src/session/prompt.ts
@@ -1243,7 +1243,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
        { message: info, parts },
      )

-      const parsed = MessageV2.Info.safeParse(info)
+      const parsed = MessageV2.Info.zod.safeParse(info)
      if (!parsed.success) {
        log.error("invalid user message before save", {
          sessionID: input.sessionID,
@@ -1254,7 +1254,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
        })
      }
      parts.forEach((part, index) => {
-        const p = MessageV2.Part.safeParse(part)
+        const p = MessageV2.Part.zod.safeParse(part)
        if (p.success) return
        log.error("invalid user part before save", {
          sessionID: input.sessionID,
--- a/packages/opencode/src/session/session.ts
+++ b/packages/opencode/src/session/session.ts
@@ -247,7 +247,7 @@ export const Event = {
    z.object({
      sessionID: SessionID.zod.optional(),
      // z.lazy defers access to break circular dep: session → message-v2 → provider → plugin → session
-      error: z.lazy(() => MessageV2.Assistant.shape.error),
+      error: z.lazy(() => (MessageV2.Assistant.zod as unknown as z.ZodObject<any>).shape.error),
    }),
  ),
 }
--- a/packages/opencode/test/session/structured-output.test.ts
+++ b/packages/opencode/test/session/structured-output.test.ts
@@ -95,7 +95,7 @@ describe("structured-output.StructuredOutputError", () => {

 describe("structured-output.UserMessage", () => {
  test("user message accepts outputFormat", () => {
-    const result = MessageV2.User.safeParse({
+    const result = MessageV2.User.zod.safeParse({
      id: MessageID.ascending(),
      sessionID: SessionID.descending(),
      role: "user",
@@ -111,7 +111,7 @@ describe("structured-output.UserMessage", () => {
  })

  test("user message works without outputFormat (optional)", () => {
-    const result = MessageV2.User.safeParse({
+    const result = MessageV2.User.zod.safeParse({
      id: MessageID.ascending(),
      sessionID: SessionID.descending(),
      role: "user",
@@ -140,7 +140,7 @@ describe("structured-output.AssistantMessage", () => {
  }

  test("assistant message accepts structured", () => {
-    const result = MessageV2.Assistant.safeParse({
+    const result = MessageV2.Assistant.zod.safeParse({
      ...baseAssistantMessage,
      structured: { company: "Anthropic", founded: 2021 },
    })
@@ -151,7 +151,7 @@ describe("structured-output.AssistantMessage", () => {
  })

  test("assistant message works without structured_output (optional)", () => {
-    const result = MessageV2.Assistant.safeParse(baseAssistantMessage)
+    const result = MessageV2.Assistant.zod.safeParse(baseAssistantMessage)
    expect(result.success).toBe(true)
  })
 })