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

6 files changed, 133 insertions, 117 deletions

diff --git a/packages/opencode/src/cli/cmd/import.ts b/packages/opencode/src/cli/cmd/import.ts
index 8da254f15..309ec6d95 100644
--- 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
diff --git a/packages/opencode/src/server/routes/instance/session.ts b/packages/opencode/src/server/routes/instance/session.ts
index 5d1f86931..8d0302426 100644
--- 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),
               },
             },
           },
diff --git a/packages/opencode/src/session/message-v2.ts b/packages/opencode/src/session/message-v2.ts
index 1a12b51eb..f1cb6db21 100644
--- 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)),
+})
+  .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 User = z.infer<typeof User>
-
 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[]
diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts
index 9d50db4af..508c72cc8 100644
--- 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,
diff --git a/packages/opencode/src/session/session.ts b/packages/opencode/src/session/session.ts
index 6e9fb5c5d..a7607798b 100644
--- 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),
     }),
   ),
 }
diff --git a/packages/opencode/test/session/structured-output.test.ts b/packages/opencode/test/session/structured-output.test.ts
index a91446bf4..c734a182a 100644
--- 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)
   })
 })