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

This commit is contained in:
Kit Langton
2026-04-21 23:26:12 -04:00
committed by GitHub
parent 1a76799fd8
commit e89543811c
6 changed files with 132 additions and 116 deletions
+2 -2
View File
@@ -168,7 +168,7 @@ export const ImportCommand = cmd({
) )
for (const msg of exportData.messages) { 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 const { id, sessionID: _, ...msgData } = msgInfo
Database.use((db) => Database.use((db) =>
db db
@@ -184,7 +184,7 @@ export const ImportCommand = cmd({
) )
for (const part of msg.parts) { 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 const { id: partId, sessionID: _s, messageID, ...partData } = partInfo
Database.use((db) => Database.use((db) =>
db db
@@ -611,7 +611,7 @@ export const SessionRoutes = lazy(() =>
description: "List of messages", description: "List of messages",
content: { content: {
"application/json": { "application/json": {
schema: resolver(MessageV2.WithParts.array()), schema: resolver(MessageV2.WithParts.zod.array()),
}, },
}, },
}, },
@@ -701,8 +701,8 @@ export const SessionRoutes = lazy(() =>
"application/json": { "application/json": {
schema: resolver( schema: resolver(
z.object({ z.object({
info: MessageV2.Info, info: MessageV2.Info.zod,
parts: MessageV2.Part.array(), parts: MessageV2.Part.zod.array(),
}), }),
), ),
}, },
@@ -813,7 +813,7 @@ export const SessionRoutes = lazy(() =>
description: "Successfully updated part", description: "Successfully updated part",
content: { content: {
"application/json": { "application/json": {
schema: resolver(MessageV2.Part), schema: resolver(MessageV2.Part.zod),
}, },
}, },
}, },
@@ -828,7 +828,7 @@ export const SessionRoutes = lazy(() =>
partID: PartID.zod, partID: PartID.zod,
}), }),
), ),
validator("json", MessageV2.Part), validator("json", MessageV2.Part.zod),
async (c) => { async (c) => {
const params = c.req.valid("param") const params = c.req.valid("param")
const body = c.req.valid("json") const body = c.req.valid("json")
@@ -856,8 +856,8 @@ export const SessionRoutes = lazy(() =>
"application/json": { "application/json": {
schema: resolver( schema: resolver(
z.object({ z.object({
info: MessageV2.Assistant, info: MessageV2.Assistant.zod,
parts: MessageV2.Part.array(), parts: MessageV2.Part.zod.array(),
}), }),
), ),
}, },
@@ -944,8 +944,8 @@ export const SessionRoutes = lazy(() =>
"application/json": { "application/json": {
schema: resolver( schema: resolver(
z.object({ z.object({
info: MessageV2.Assistant, info: MessageV2.Assistant.zod,
parts: MessageV2.Part.array(), parts: MessageV2.Part.zod.array(),
}), }),
), ),
}, },
@@ -980,7 +980,7 @@ export const SessionRoutes = lazy(() =>
description: "Created message", description: "Created message",
content: { content: {
"application/json": { "application/json": {
schema: resolver(MessageV2.WithParts), schema: resolver(MessageV2.WithParts.zod),
}, },
}, },
}, },
+113 -97
View File
@@ -368,37 +368,68 @@ export type ToolPart = Omit<Types.DeepMutable<Schema.Schema.Type<typeof ToolPart
state: ToolState state: ToolState
} }
const Base = z.object({ const messageBase = {
id: MessageID.zod, id: MessageID,
sessionID: SessionID.zod, sessionID: SessionID,
}) }
export const User = Base.extend({ export const User = Schema.Struct({
role: z.literal("user"), ...messageBase,
time: z.object({ role: Schema.Literal("user"),
created: z.number(), time: Schema.Struct({
created: Schema.Number,
}), }),
format: Format.zod.optional(), format: Schema.optional(_Format),
summary: z summary: Schema.optional(
.object({ Schema.Struct({
title: z.string().optional(), title: Schema.optional(Schema.String),
body: z.string().optional(), body: Schema.optional(Schema.String),
diffs: Snapshot.FileDiff.zod.array(), diffs: Schema.Array(Snapshot.FileDiff),
}) }),
.optional(), ),
agent: z.string(), agent: Schema.String,
model: z.object({ model: Schema.Struct({
providerID: ProviderID.zod, providerID: ProviderID,
modelID: ModelID.zod, modelID: ModelID,
variant: z.string().optional(), variant: Schema.optional(Schema.String),
}), }),
system: z.string().optional(), system: Schema.optional(Schema.String),
tools: z.record(z.string(), z.boolean()).optional(), tools: Schema.optional(Schema.Record(Schema.String, Schema.Boolean)),
}).meta({
ref: "UserMessage",
}) })
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 = export type Part =
| TextPart | TextPart
| SubtaskPart | SubtaskPart
@@ -413,28 +444,19 @@ export type Part =
| RetryPart | RetryPart
| CompactionPart | CompactionPart
// The derived `.zod` on each leaf is typed as `z.ZodType<...>`, but the walker // Errors are still NamedError-based Zod; bridge via ZodOverride so the derived
// always emits a `z.ZodObject` at runtime. `z.discriminatedUnion` and // Zod + JSON Schema emit the original discriminatedUnion shape. Migrating the
// `z.infer` both rely on the ZodObject structural type, so cast here so the // error classes to Schema.TaggedErrorClass is a separate slice.
// resulting Part behaves like the pre-migration Zod union. const AssistantErrorZod = z.discriminatedUnion("name", [
export const Part = z AuthError.Schema,
.discriminatedUnion("type", [ NamedError.Unknown.Schema,
TextPart.zod as unknown as z.ZodObject<any>, OutputLengthError.Schema,
SubtaskPart.zod as unknown as z.ZodObject<any>, AbortedError.Schema,
ReasoningPart.zod as unknown as z.ZodObject<any>, StructuredOutputError.Schema,
FilePart.zod as unknown as z.ZodObject<any>, ContextOverflowError.Schema,
ToolPart.zod as unknown as z.ZodObject<any>, APIError.Schema,
StepStartPart.zod as unknown as z.ZodObject<any>, ])
StepFinishPart.zod as unknown as z.ZodObject<any>, type AssistantError = z.infer<typeof AssistantErrorZod>
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>
// ── Prompt input schemas ───────────────────────────────────────────────────── // ── Prompt input schemas ─────────────────────────────────────────────────────
// //
@@ -508,59 +530,53 @@ export const SubtaskPartInput = Schema.Struct({
.pipe(withStatics((s) => ({ zod: zod(s) }))) .pipe(withStatics((s) => ({ zod: zod(s) })))
export type SubtaskPartInput = Types.DeepMutable<Schema.Schema.Type<typeof SubtaskPartInput>> export type SubtaskPartInput = Types.DeepMutable<Schema.Schema.Type<typeof SubtaskPartInput>>
export const Assistant = Base.extend({ export const Assistant = Schema.Struct({
role: z.literal("assistant"), ...messageBase,
time: z.object({ role: Schema.Literal("assistant"),
created: z.number(), time: Schema.Struct({
completed: z.number().optional(), created: Schema.Number,
completed: Schema.optional(Schema.Number),
}), }),
error: z error: Schema.optional(Schema.Any.annotate({ [ZodOverride]: AssistantErrorZod })),
.discriminatedUnion("name", [ parentID: MessageID,
AuthError.Schema, modelID: ModelID,
NamedError.Unknown.Schema, providerID: ProviderID,
OutputLengthError.Schema,
AbortedError.Schema,
StructuredOutputError.Schema,
ContextOverflowError.Schema,
APIError.Schema,
])
.optional(),
parentID: MessageID.zod,
modelID: ModelID.zod,
providerID: ProviderID.zod,
/** /**
* @deprecated * @deprecated
*/ */
mode: z.string(), mode: Schema.String,
agent: z.string(), agent: Schema.String,
path: z.object({ path: Schema.Struct({
cwd: z.string(), cwd: Schema.String,
root: z.string(), root: Schema.String,
}), }),
summary: z.boolean().optional(), summary: Schema.optional(Schema.Boolean),
cost: z.number(), cost: Schema.Number,
tokens: z.object({ tokens: Schema.Struct({
total: z.number().optional(), total: Schema.optional(Schema.Number),
input: z.number(), input: Schema.Number,
output: z.number(), output: Schema.Number,
reasoning: z.number(), reasoning: Schema.Number,
cache: z.object({ cache: Schema.Struct({
read: z.number(), read: Schema.Number,
write: z.number(), write: Schema.Number,
}), }),
}), }),
structured: z.any().optional(), structured: Schema.optional(Schema.Any),
variant: z.string().optional(), variant: Schema.optional(Schema.String),
finish: z.string().optional(), finish: Schema.optional(Schema.String),
}).meta({
ref: "AssistantMessage",
}) })
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({ const _Info = Schema.Union([User, Assistant]).annotate({ discriminator: "role", identifier: "Message" })
ref: "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 = { export const Event = {
Updated: SyncEvent.define({ Updated: SyncEvent.define({
@@ -569,7 +585,7 @@ export const Event = {
aggregate: "sessionID", aggregate: "sessionID",
schema: z.object({ schema: z.object({
sessionID: SessionID.zod, sessionID: SessionID.zod,
info: Info, info: Info.zod,
}), }),
}), }),
Removed: SyncEvent.define({ Removed: SyncEvent.define({
@@ -587,7 +603,7 @@ export const Event = {
aggregate: "sessionID", aggregate: "sessionID",
schema: z.object({ schema: z.object({
sessionID: SessionID.zod, sessionID: SessionID.zod,
part: Part, part: Part.zod,
time: z.number(), time: z.number(),
}), }),
}), }),
@@ -613,10 +629,10 @@ export const Event = {
}), }),
} }
export const WithParts = z.object({ export const WithParts = Schema.Struct({
info: Info, info: _Info,
parts: z.array(Part), parts: Schema.Array(_Part),
}) }).pipe(withStatics((s) => ({ zod: zod(s) })))
export type WithParts = { export type WithParts = {
info: Info info: Info
parts: Part[] parts: Part[]
+2 -2
View File
@@ -1243,7 +1243,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the
{ message: info, parts }, { message: info, parts },
) )
const parsed = MessageV2.Info.safeParse(info) const parsed = MessageV2.Info.zod.safeParse(info)
if (!parsed.success) { if (!parsed.success) {
log.error("invalid user message before save", { log.error("invalid user message before save", {
sessionID: input.sessionID, 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) => { parts.forEach((part, index) => {
const p = MessageV2.Part.safeParse(part) const p = MessageV2.Part.zod.safeParse(part)
if (p.success) return if (p.success) return
log.error("invalid user part before save", { log.error("invalid user part before save", {
sessionID: input.sessionID, sessionID: input.sessionID,
+1 -1
View File
@@ -247,7 +247,7 @@ export const Event = {
z.object({ z.object({
sessionID: SessionID.zod.optional(), sessionID: SessionID.zod.optional(),
// z.lazy defers access to break circular dep: session → message-v2 → provider → plugin → session // 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),
}), }),
), ),
} }
@@ -95,7 +95,7 @@ describe("structured-output.StructuredOutputError", () => {
describe("structured-output.UserMessage", () => { describe("structured-output.UserMessage", () => {
test("user message accepts outputFormat", () => { test("user message accepts outputFormat", () => {
const result = MessageV2.User.safeParse({ const result = MessageV2.User.zod.safeParse({
id: MessageID.ascending(), id: MessageID.ascending(),
sessionID: SessionID.descending(), sessionID: SessionID.descending(),
role: "user", role: "user",
@@ -111,7 +111,7 @@ describe("structured-output.UserMessage", () => {
}) })
test("user message works without outputFormat (optional)", () => { test("user message works without outputFormat (optional)", () => {
const result = MessageV2.User.safeParse({ const result = MessageV2.User.zod.safeParse({
id: MessageID.ascending(), id: MessageID.ascending(),
sessionID: SessionID.descending(), sessionID: SessionID.descending(),
role: "user", role: "user",
@@ -140,7 +140,7 @@ describe("structured-output.AssistantMessage", () => {
} }
test("assistant message accepts structured", () => { test("assistant message accepts structured", () => {
const result = MessageV2.Assistant.safeParse({ const result = MessageV2.Assistant.zod.safeParse({
...baseAssistantMessage, ...baseAssistantMessage,
structured: { company: "Anthropic", founded: 2021 }, structured: { company: "Anthropic", founded: 2021 },
}) })
@@ -151,7 +151,7 @@ describe("structured-output.AssistantMessage", () => {
}) })
test("assistant message works without structured_output (optional)", () => { 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) expect(result.success).toBe(true)
}) })
}) })