refactor(core): migrate MessageV2 tool state schemas to Effect Schema (#23752)

1 files changed, 74 insertions, 68 deletions

diff --git a/packages/opencode/src/session/message-v2.ts b/packages/opencode/src/session/message-v2.ts
index 83477d12b..04cb15ef8 100644
--- a/packages/opencode/src/session/message-v2.ts
+++ b/packages/opencode/src/session/message-v2.ts
@@ -15,8 +15,9 @@ import { isMedia } from "@/util/media"
 import type { SystemError } from "bun"
 import type { Provider } from "@/provider"
 import { ModelID, ProviderID } from "@/provider/schema"
-import { Effect, Schema } from "effect"
-import { zod } from "@/util/effect-zod"
+import { Effect, Schema, Types } from "effect"
+import { zod, ZodOverride } from "@/util/effect-zod"
+import { withStatics } from "@/util/schema"
 import { EffectLogger } from "@/effect"
 
 /** Error shape thrown by Bun's fetch() when gzip/br decompression fails mid-stream */
@@ -272,79 +273,84 @@ export const StepFinishPart = PartBase.extend({
 })
 export type StepFinishPart = z.infer<typeof StepFinishPart>
 
-export const ToolStatePending = z
-  .object({
-    status: z.literal("pending"),
-    input: z.record(z.string(), z.any()),
-    raw: z.string(),
-  })
-  .meta({
-    ref: "ToolStatePending",
-  })
-
-export type ToolStatePending = z.infer<typeof ToolStatePending>
+export const ToolStatePending = Schema.Struct({
+  status: Schema.Literal("pending"),
+  input: Schema.Record(Schema.String, Schema.Any),
+  raw: Schema.String,
+})
+  .annotate({ identifier: "ToolStatePending" })
+  .pipe(withStatics((s) => ({ zod: zod(s) })))
+export type ToolStatePending = Types.DeepMutable<Schema.Schema.Type<typeof ToolStatePending>>
+
+export const ToolStateRunning = Schema.Struct({
+  status: Schema.Literal("running"),
+  input: Schema.Record(Schema.String, Schema.Any),
+  title: Schema.optional(Schema.String),
+  metadata: Schema.optional(Schema.Record(Schema.String, Schema.Any)),
+  time: Schema.Struct({
+    start: Schema.Number,
+  }),
+})
+  .annotate({ identifier: "ToolStateRunning" })
+  .pipe(withStatics((s) => ({ zod: zod(s) })))
+export type ToolStateRunning = Types.DeepMutable<Schema.Schema.Type<typeof ToolStateRunning>>
+
+export const ToolStateCompleted = Schema.Struct({
+  status: Schema.Literal("completed"),
+  input: Schema.Record(Schema.String, Schema.Any),
+  output: Schema.String,
+  title: Schema.String,
+  metadata: Schema.Record(Schema.String, Schema.Any),
+  time: Schema.Struct({
+    start: Schema.Number,
+    end: Schema.Number,
+    compacted: Schema.optional(Schema.Number),
+  }),
+  // FilePart is still Zod-first this slice; bridge via ZodOverride so the
+  // derived Zod + JSON Schema still emit `$ref: FilePart` array items.
+  attachments: Schema.optional(Schema.Any.annotate({ [ZodOverride]: FilePart.array() })),
+})
+  .annotate({ identifier: "ToolStateCompleted" })
+  .pipe(withStatics((s) => ({ zod: zod(s) })))
+export type ToolStateCompleted = Omit<
+  Types.DeepMutable<Schema.Schema.Type<typeof ToolStateCompleted>>,
+  "attachments"
+> & {
+  attachments?: FilePart[]
+}
 
-export const ToolStateRunning = z
-  .object({
-    status: z.literal("running"),
-    input: z.record(z.string(), z.any()),
-    title: z.string().optional(),
-    metadata: z.record(z.string(), z.any()).optional(),
-    time: z.object({
-      start: z.number(),
-    }),
-  })
-  .meta({
-    ref: "ToolStateRunning",
-  })
-export type ToolStateRunning = z.infer<typeof ToolStateRunning>
-
-export const ToolStateCompleted = z
-  .object({
-    status: z.literal("completed"),
-    input: z.record(z.string(), z.any()),
-    output: z.string(),
-    title: z.string(),
-    metadata: z.record(z.string(), z.any()),
-    time: z.object({
-      start: z.number(),
-      end: z.number(),
-      compacted: z.number().optional(),
-    }),
-    attachments: FilePart.array().optional(),
-  })
-  .meta({
-    ref: "ToolStateCompleted",
-  })
-export type ToolStateCompleted = z.infer<typeof ToolStateCompleted>
-
-export const ToolStateError = z
-  .object({
-    status: z.literal("error"),
-    input: z.record(z.string(), z.any()),
-    error: z.string(),
-    metadata: z.record(z.string(), z.any()).optional(),
-    time: z.object({
-      start: z.number(),
-      end: z.number(),
-    }),
-  })
-  .meta({
-    ref: "ToolStateError",
-  })
-export type ToolStateError = z.infer<typeof ToolStateError>
+export const ToolStateError = Schema.Struct({
+  status: Schema.Literal("error"),
+  input: Schema.Record(Schema.String, Schema.Any),
+  error: Schema.String,
+  metadata: Schema.optional(Schema.Record(Schema.String, Schema.Any)),
+  time: Schema.Struct({
+    start: Schema.Number,
+    end: Schema.Number,
+  }),
+})
+  .annotate({ identifier: "ToolStateError" })
+  .pipe(withStatics((s) => ({ zod: zod(s) })))
+export type ToolStateError = Types.DeepMutable<Schema.Schema.Type<typeof ToolStateError>>
 
-export const ToolState = z
-  .discriminatedUnion("status", [ToolStatePending, ToolStateRunning, ToolStateCompleted, ToolStateError])
-  .meta({
-    ref: "ToolState",
-  })
+const _ToolState = Schema.Union([ToolStatePending, ToolStateRunning, ToolStateCompleted, ToolStateError]).annotate({
+  discriminator: "status",
+  identifier: "ToolState",
+})
+// Cast the derived zod so downstream z.infer sees the same mutable shape that
+// our exported TS types expose (the pre-migration Zod inferences were mutable).
+export const ToolState = Object.assign(_ToolState, {
+  zod: zod(_ToolState) as unknown as z.ZodType<
+    ToolStatePending | ToolStateRunning | ToolStateCompleted | ToolStateError
+  >,
+})
+export type ToolState = ToolStatePending | ToolStateRunning | ToolStateCompleted | ToolStateError
 
 export const ToolPart = PartBase.extend({
   type: z.literal("tool"),
   callID: z.string(),
   tool: z.string(),
-  state: ToolState,
+  state: ToolState.zod,
   metadata: z.record(z.string(), z.any()).optional(),
 }).meta({
   ref: "ToolPart",