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

This commit is contained in:
Kit Langton
2026-04-21 17:47:50 -04:00
committed by GitHub
parent d6dea3f3e0
commit df0c1f649c
+64 -58
View File
@@ -15,8 +15,9 @@ import { isMedia } from "@/util/media"
import type { SystemError } from "bun" import type { SystemError } from "bun"
import type { Provider } from "@/provider" import type { Provider } from "@/provider"
import { ModelID, ProviderID } from "@/provider/schema" import { ModelID, ProviderID } from "@/provider/schema"
import { Effect, Schema } from "effect" import { Effect, Schema, Types } from "effect"
import { zod } from "@/util/effect-zod" import { zod, ZodOverride } from "@/util/effect-zod"
import { withStatics } from "@/util/schema"
import { EffectLogger } from "@/effect" import { EffectLogger } from "@/effect"
/** Error shape thrown by Bun's fetch() when gzip/br decompression fails mid-stream */ /** 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 type StepFinishPart = z.infer<typeof StepFinishPart>
export const ToolStatePending = z export const ToolStatePending = Schema.Struct({
.object({ status: Schema.Literal("pending"),
status: z.literal("pending"), input: Schema.Record(Schema.String, Schema.Any),
input: z.record(z.string(), z.any()), raw: Schema.String,
raw: z.string(),
})
.meta({
ref: "ToolStatePending",
}) })
.annotate({ identifier: "ToolStatePending" })
.pipe(withStatics((s) => ({ zod: zod(s) })))
export type ToolStatePending = Types.DeepMutable<Schema.Schema.Type<typeof ToolStatePending>>
export type ToolStatePending = z.infer<typeof ToolStatePending> export const ToolStateRunning = Schema.Struct({
status: Schema.Literal("running"),
export const ToolStateRunning = z input: Schema.Record(Schema.String, Schema.Any),
.object({ title: Schema.optional(Schema.String),
status: z.literal("running"), metadata: Schema.optional(Schema.Record(Schema.String, Schema.Any)),
input: z.record(z.string(), z.any()), time: Schema.Struct({
title: z.string().optional(), start: Schema.Number,
metadata: z.record(z.string(), z.any()).optional(),
time: z.object({
start: z.number(),
}), }),
}) })
.meta({ .annotate({ identifier: "ToolStateRunning" })
ref: "ToolStateRunning", .pipe(withStatics((s) => ({ zod: zod(s) })))
}) export type ToolStateRunning = Types.DeepMutable<Schema.Schema.Type<typeof ToolStateRunning>>
export type ToolStateRunning = z.infer<typeof ToolStateRunning>
export const ToolStateCompleted = z export const ToolStateCompleted = Schema.Struct({
.object({ status: Schema.Literal("completed"),
status: z.literal("completed"), input: Schema.Record(Schema.String, Schema.Any),
input: z.record(z.string(), z.any()), output: Schema.String,
output: z.string(), title: Schema.String,
title: z.string(), metadata: Schema.Record(Schema.String, Schema.Any),
metadata: z.record(z.string(), z.any()), time: Schema.Struct({
time: z.object({ start: Schema.Number,
start: z.number(), end: Schema.Number,
end: z.number(), compacted: Schema.optional(Schema.Number),
compacted: z.number().optional(),
}), }),
attachments: FilePart.array().optional(), // 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() })),
}) })
.meta({ .annotate({ identifier: "ToolStateCompleted" })
ref: "ToolStateCompleted", .pipe(withStatics((s) => ({ zod: zod(s) })))
}) export type ToolStateCompleted = Omit<
export type ToolStateCompleted = z.infer<typeof ToolStateCompleted> Types.DeepMutable<Schema.Schema.Type<typeof ToolStateCompleted>>,
"attachments"
> & {
attachments?: FilePart[]
}
export const ToolStateError = z export const ToolStateError = Schema.Struct({
.object({ status: Schema.Literal("error"),
status: z.literal("error"), input: Schema.Record(Schema.String, Schema.Any),
input: z.record(z.string(), z.any()), error: Schema.String,
error: z.string(), metadata: Schema.optional(Schema.Record(Schema.String, Schema.Any)),
metadata: z.record(z.string(), z.any()).optional(), time: Schema.Struct({
time: z.object({ start: Schema.Number,
start: z.number(), end: Schema.Number,
end: z.number(),
}), }),
}) })
.meta({ .annotate({ identifier: "ToolStateError" })
ref: "ToolStateError", .pipe(withStatics((s) => ({ zod: zod(s) })))
}) export type ToolStateError = Types.DeepMutable<Schema.Schema.Type<typeof ToolStateError>>
export type ToolStateError = z.infer<typeof ToolStateError>
export const ToolState = z const _ToolState = Schema.Union([ToolStatePending, ToolStateRunning, ToolStateCompleted, ToolStateError]).annotate({
.discriminatedUnion("status", [ToolStatePending, ToolStateRunning, ToolStateCompleted, ToolStateError]) discriminator: "status",
.meta({ identifier: "ToolState",
ref: "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({ export const ToolPart = PartBase.extend({
type: z.literal("tool"), type: z.literal("tool"),
callID: z.string(), callID: z.string(),
tool: z.string(), tool: z.string(),
state: ToolState, state: ToolState.zod,
metadata: z.record(z.string(), z.any()).optional(), metadata: z.record(z.string(), z.any()).optional(),
}).meta({ }).meta({
ref: "ToolPart", ref: "ToolPart",