From e89543811ca02067732b9ae6637bc1c1572dc7c1 Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Tue, 21 Apr 2026 23:26:12 -0400 Subject: [PATCH] refactor(core): migrate MessageV2 message DTOs (User/Assistant/Part/Info/WithParts) to Effect Schema (#23757) --- packages/opencode/src/cli/cmd/import.ts | 4 +- .../src/server/routes/instance/session.ts | 20 +- packages/opencode/src/session/message-v2.ts | 210 ++++++++++-------- packages/opencode/src/session/prompt.ts | 4 +- packages/opencode/src/session/session.ts | 2 +- .../test/session/structured-output.test.ts | 8 +- 6 files changed, 132 insertions(+), 116 deletions(-) diff --git a/packages/opencode/src/cli/cmd/import.ts b/packages/opencode/src/cli/cmd/import.ts index 8da254f159..309ec6d950 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 5d1f869310..8d03024260 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 1a12b51eb8..f1cb6db218 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 + .annotate({ identifier: "UserMessage" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type User = Types.DeepMutable> +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, - SubtaskPart.zod as unknown as z.ZodObject, - ReasoningPart.zod as unknown as z.ZodObject, - FilePart.zod as unknown as z.ZodObject, - ToolPart.zod as unknown as z.ZodObject, - StepStartPart.zod as unknown as z.ZodObject, - StepFinishPart.zod as unknown as z.ZodObject, - SnapshotPart.zod as unknown as z.ZodObject, - PatchPart.zod as unknown as z.ZodObject, - AgentPart.zod as unknown as z.ZodObject, - RetryPart.zod as unknown as z.ZodObject, - CompactionPart.zod as unknown as z.ZodObject, - ]) - .meta({ - ref: "Part", - }) as unknown as z.ZodType +// 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 // ── Prompt input schemas ───────────────────────────────────────────────────── // @@ -508,59 +530,53 @@ export const SubtaskPartInput = Schema.Struct({ .pipe(withStatics((s) => ({ zod: zod(s) }))) export type SubtaskPartInput = Types.DeepMutable> -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 + .annotate({ identifier: "AssistantMessage" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type Assistant = Omit>, "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, }) -export type Info = z.infer +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 9d50db4afb..508c72cc8f 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 6e9fb5c5d8..a7607798ba 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).shape.error), }), ), } diff --git a/packages/opencode/test/session/structured-output.test.ts b/packages/opencode/test/session/structured-output.test.ts index a91446bf42..c734a182ae 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) }) })