From fa623964a262958f1225afef1e4b286b682ea19f Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Tue, 21 Apr 2026 23:17:23 -0400 Subject: [PATCH] refactor(core): migrate MessageV2 part leaves + ToolPart to Effect Schema (#23756) --- .../src/server/routes/instance/session.ts | 6 +- packages/opencode/src/session/message-v2.ts | 509 +++++++++++------- packages/opencode/src/session/prompt.ts | 72 +-- 3 files changed, 336 insertions(+), 251 deletions(-) diff --git a/packages/opencode/src/server/routes/instance/session.ts b/packages/opencode/src/server/routes/instance/session.ts index a46c2f3bf3..adafb8f360 100644 --- a/packages/opencode/src/server/routes/instance/session.ts +++ b/packages/opencode/src/server/routes/instance/session.ts @@ -882,7 +882,9 @@ export const SessionRoutes = lazy(() => const msg = await runRequest( "SessionRoutes.prompt", c, - SessionPrompt.Service.use((svc) => svc.prompt({ ...body, sessionID })), + SessionPrompt.Service.use((svc) => + svc.prompt({ ...body, sessionID } as unknown as SessionPrompt.PromptInput), + ), ) void stream.write(JSON.stringify(msg)) }) @@ -915,7 +917,7 @@ export const SessionRoutes = lazy(() => void runRequest( "SessionRoutes.prompt_async", c, - SessionPrompt.Service.use((svc) => svc.prompt({ ...body, sessionID })), + SessionPrompt.Service.use((svc) => svc.prompt({ ...body, sessionID } as unknown as SessionPrompt.PromptInput)), ).catch((err) => { log.error("prompt_async failed", { sessionID, error: err }) void Bus.publish(Session.Event.Error, { diff --git a/packages/opencode/src/session/message-v2.ts b/packages/opencode/src/session/message-v2.ts index 04cb15ef8d..1a12b51eb8 100644 --- a/packages/opencode/src/session/message-v2.ts +++ b/packages/opencode/src/session/message-v2.ts @@ -86,192 +86,207 @@ const _Format = Schema.Union([OutputFormatText, OutputFormatJsonSchema]).annotat export const Format = Object.assign(_Format, { zod: zod(_Format) }) export type OutputFormat = Schema.Schema.Type -const PartBase = z.object({ - id: PartID.zod, - sessionID: SessionID.zod, - messageID: MessageID.zod, -}) +const partBase = { + id: PartID, + sessionID: SessionID, + messageID: MessageID, +} -export const SnapshotPart = PartBase.extend({ - type: z.literal("snapshot"), - snapshot: z.string(), -}).meta({ - ref: "SnapshotPart", +export const SnapshotPart = Schema.Struct({ + ...partBase, + type: Schema.Literal("snapshot"), + snapshot: Schema.String, }) -export type SnapshotPart = z.infer + .annotate({ identifier: "SnapshotPart" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type SnapshotPart = Types.DeepMutable> -export const PatchPart = PartBase.extend({ - type: z.literal("patch"), - hash: z.string(), - files: z.string().array(), -}).meta({ - ref: "PatchPart", +export const PatchPart = Schema.Struct({ + ...partBase, + type: Schema.Literal("patch"), + hash: Schema.String, + files: Schema.Array(Schema.String), }) -export type PatchPart = z.infer + .annotate({ identifier: "PatchPart" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type PatchPart = Types.DeepMutable> -export const TextPart = PartBase.extend({ - type: z.literal("text"), - text: z.string(), - synthetic: z.boolean().optional(), - ignored: z.boolean().optional(), - time: z - .object({ - start: z.number(), - end: z.number().optional(), - }) - .optional(), - metadata: z.record(z.string(), z.any()).optional(), -}).meta({ - ref: "TextPart", -}) -export type TextPart = z.infer - -export const ReasoningPart = PartBase.extend({ - type: z.literal("reasoning"), - text: z.string(), - metadata: z.record(z.string(), z.any()).optional(), - time: z.object({ - start: z.number(), - end: z.number().optional(), - }), -}).meta({ - ref: "ReasoningPart", -}) -export type ReasoningPart = z.infer - -const FilePartSourceBase = z.object({ - text: z - .object({ - value: z.string(), - start: z.number().int(), - end: z.number().int(), - }) - .meta({ - ref: "FilePartSourceText", +export const TextPart = Schema.Struct({ + ...partBase, + type: Schema.Literal("text"), + text: Schema.String, + synthetic: Schema.optional(Schema.Boolean), + ignored: Schema.optional(Schema.Boolean), + time: Schema.optional( + Schema.Struct({ + start: Schema.Number, + end: Schema.optional(Schema.Number), }), + ), + metadata: Schema.optional(Schema.Record(Schema.String, Schema.Any)), }) + .annotate({ identifier: "TextPart" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type TextPart = Types.DeepMutable> -export const FileSource = FilePartSourceBase.extend({ - type: z.literal("file"), - path: z.string(), -}).meta({ - ref: "FileSource", -}) - -export const SymbolSource = FilePartSourceBase.extend({ - type: z.literal("symbol"), - path: z.string(), - range: LSP.Range.zod, - name: z.string(), - kind: z.number().int(), -}).meta({ - ref: "SymbolSource", -}) - -export const ResourceSource = FilePartSourceBase.extend({ - type: z.literal("resource"), - clientName: z.string(), - uri: z.string(), -}).meta({ - ref: "ResourceSource", -}) - -export const FilePartSource = z.discriminatedUnion("type", [FileSource, SymbolSource, ResourceSource]).meta({ - ref: "FilePartSource", -}) - -export const FilePart = PartBase.extend({ - type: z.literal("file"), - mime: z.string(), - filename: z.string().optional(), - url: z.string(), - source: FilePartSource.optional(), -}).meta({ - ref: "FilePart", -}) -export type FilePart = z.infer - -export const AgentPart = PartBase.extend({ - type: z.literal("agent"), - name: z.string(), - source: z - .object({ - value: z.string(), - start: z.number().int(), - end: z.number().int(), - }) - .optional(), -}).meta({ - ref: "AgentPart", -}) -export type AgentPart = z.infer - -export const CompactionPart = PartBase.extend({ - type: z.literal("compaction"), - auto: z.boolean(), - overflow: z.boolean().optional(), - tail_start_id: MessageID.zod.optional(), -}).meta({ - ref: "CompactionPart", -}) -export type CompactionPart = z.infer - -export const SubtaskPart = PartBase.extend({ - type: z.literal("subtask"), - prompt: z.string(), - description: z.string(), - agent: z.string(), - model: z - .object({ - providerID: ProviderID.zod, - modelID: ModelID.zod, - }) - .optional(), - command: z.string().optional(), -}).meta({ - ref: "SubtaskPart", -}) -export type SubtaskPart = z.infer - -export const RetryPart = PartBase.extend({ - type: z.literal("retry"), - attempt: z.number(), - error: APIError.Schema, - time: z.object({ - created: z.number(), +export const ReasoningPart = Schema.Struct({ + ...partBase, + type: Schema.Literal("reasoning"), + text: Schema.String, + metadata: Schema.optional(Schema.Record(Schema.String, Schema.Any)), + time: Schema.Struct({ + start: Schema.Number, + end: Schema.optional(Schema.Number), }), -}).meta({ - ref: "RetryPart", }) -export type RetryPart = z.infer + .annotate({ identifier: "ReasoningPart" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type ReasoningPart = Types.DeepMutable> -export const StepStartPart = PartBase.extend({ - type: z.literal("step-start"), - snapshot: z.string().optional(), -}).meta({ - ref: "StepStartPart", +const filePartSourceBase = { + text: Schema.Struct({ + value: Schema.String, + start: Schema.Number.check(Schema.isInt()), + end: Schema.Number.check(Schema.isInt()), + }).annotate({ identifier: "FilePartSourceText" }), +} + +export const FileSource = Schema.Struct({ + ...filePartSourceBase, + type: Schema.Literal("file"), + path: Schema.String, }) -export type StepStartPart = z.infer + .annotate({ identifier: "FileSource" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) -export const StepFinishPart = PartBase.extend({ - type: z.literal("step-finish"), - reason: z.string(), - snapshot: z.string().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(), +export const SymbolSource = Schema.Struct({ + ...filePartSourceBase, + type: Schema.Literal("symbol"), + path: Schema.String, + range: LSP.Range, + name: Schema.String, + kind: Schema.Number.check(Schema.isInt()), +}) + .annotate({ identifier: "SymbolSource" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) + +export const ResourceSource = Schema.Struct({ + ...filePartSourceBase, + type: Schema.Literal("resource"), + clientName: Schema.String, + uri: Schema.String, +}) + .annotate({ identifier: "ResourceSource" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) + +const _FilePartSource = Schema.Union([FileSource, SymbolSource, ResourceSource]).annotate({ + discriminator: "type", + identifier: "FilePartSource", +}) +export const FilePartSource = Object.assign(_FilePartSource, { zod: zod(_FilePartSource) }) + +export const FilePart = Schema.Struct({ + ...partBase, + type: Schema.Literal("file"), + mime: Schema.String, + filename: Schema.optional(Schema.String), + url: Schema.String, + source: Schema.optional(_FilePartSource), +}) + .annotate({ identifier: "FilePart" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type FilePart = Types.DeepMutable> + +export const AgentPart = Schema.Struct({ + ...partBase, + type: Schema.Literal("agent"), + name: Schema.String, + source: Schema.optional( + Schema.Struct({ + value: Schema.String, + start: Schema.Number.check(Schema.isInt()), + end: Schema.Number.check(Schema.isInt()), + }), + ), +}) + .annotate({ identifier: "AgentPart" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type AgentPart = Types.DeepMutable> + +export const CompactionPart = Schema.Struct({ + ...partBase, + type: Schema.Literal("compaction"), + auto: Schema.Boolean, + overflow: Schema.optional(Schema.Boolean), + tail_start_id: Schema.optional(MessageID), +}) + .annotate({ identifier: "CompactionPart" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type CompactionPart = Types.DeepMutable> + +export const SubtaskPart = Schema.Struct({ + ...partBase, + type: Schema.Literal("subtask"), + prompt: Schema.String, + description: Schema.String, + agent: Schema.String, + model: Schema.optional( + Schema.Struct({ + providerID: ProviderID, + modelID: ModelID, + }), + ), + command: Schema.optional(Schema.String), +}) + .annotate({ identifier: "SubtaskPart" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type SubtaskPart = Types.DeepMutable> + +export const RetryPart = Schema.Struct({ + ...partBase, + type: Schema.Literal("retry"), + attempt: Schema.Number, + // APIError is still NamedError-based Zod; bridge via ZodOverride until errors migrate. + error: Schema.Any.annotate({ [ZodOverride]: APIError.Schema }), + time: Schema.Struct({ + created: Schema.Number, + }), +}) + .annotate({ identifier: "RetryPart" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type RetryPart = Omit>, "error"> & { + error: APIError +} + +export const StepStartPart = Schema.Struct({ + ...partBase, + type: Schema.Literal("step-start"), + snapshot: Schema.optional(Schema.String), +}) + .annotate({ identifier: "StepStartPart" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type StepStartPart = Types.DeepMutable> + +export const StepFinishPart = Schema.Struct({ + ...partBase, + type: Schema.Literal("step-finish"), + reason: Schema.String, + snapshot: Schema.optional(Schema.String), + 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, }), }), -}).meta({ - ref: "StepFinishPart", }) -export type StepFinishPart = z.infer + .annotate({ identifier: "StepFinishPart" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type StepFinishPart = Types.DeepMutable> export const ToolStatePending = Schema.Struct({ status: Schema.Literal("pending"), @@ -306,18 +321,11 @@ export const ToolStateCompleted = Schema.Struct({ 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() })), + attachments: Schema.optional(Schema.Array(FilePart)), }) .annotate({ identifier: "ToolStateCompleted" }) .pipe(withStatics((s) => ({ zod: zod(s) }))) -export type ToolStateCompleted = Omit< - Types.DeepMutable>, - "attachments" -> & { - attachments?: FilePart[] -} +export type ToolStateCompleted = Types.DeepMutable> export const ToolStateError = Schema.Struct({ status: Schema.Literal("error"), @@ -346,16 +354,19 @@ export const ToolState = Object.assign(_ToolState, { }) export type ToolState = ToolStatePending | ToolStateRunning | ToolStateCompleted | ToolStateError -export const ToolPart = PartBase.extend({ - type: z.literal("tool"), - callID: z.string(), - tool: z.string(), - state: ToolState.zod, - metadata: z.record(z.string(), z.any()).optional(), -}).meta({ - ref: "ToolPart", +export const ToolPart = Schema.Struct({ + ...partBase, + type: Schema.Literal("tool"), + callID: Schema.String, + tool: Schema.String, + state: _ToolState, + metadata: Schema.optional(Schema.Record(Schema.String, Schema.Any)), }) -export type ToolPart = z.infer + .annotate({ identifier: "ToolPart" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type ToolPart = Omit>, "state"> & { + state: ToolState +} const Base = z.object({ id: MessageID.zod, @@ -388,25 +399,114 @@ export const User = Base.extend({ }) export type User = z.infer +export type Part = + | TextPart + | SubtaskPart + | ReasoningPart + | FilePart + | ToolPart + | StepStartPart + | StepFinishPart + | SnapshotPart + | PatchPart + | AgentPart + | 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, - SubtaskPart, - ReasoningPart, - FilePart, - ToolPart, - StepStartPart, - StepFinishPart, - SnapshotPart, - PatchPart, - AgentPart, - RetryPart, - CompactionPart, + 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", - }) -export type Part = z.infer + }) as unknown as z.ZodType + +// ── Prompt input schemas ───────────────────────────────────────────────────── +// +// Consumers of `SessionPrompt.PromptInput.parts` send part drafts without the +// ambient IDs (`messageID`, `sessionID`) that live on stored parts, and may +// omit `id` to let the server allocate one. These Schema-Struct variants +// carry that shape, and `SessionPrompt.PromptInput` just references the +// derived `.zod` (no omit/partial gymnastics needed at the call site). + +export const TextPartInput = Schema.Struct({ + id: Schema.optional(PartID), + type: Schema.Literal("text"), + text: Schema.String, + synthetic: Schema.optional(Schema.Boolean), + ignored: Schema.optional(Schema.Boolean), + time: Schema.optional( + Schema.Struct({ + start: Schema.Number, + end: Schema.optional(Schema.Number), + }), + ), + metadata: Schema.optional(Schema.Record(Schema.String, Schema.Any)), +}) + .annotate({ identifier: "TextPartInput" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type TextPartInput = Types.DeepMutable> + +export const FilePartInput = Schema.Struct({ + id: Schema.optional(PartID), + type: Schema.Literal("file"), + mime: Schema.String, + filename: Schema.optional(Schema.String), + url: Schema.String, + source: Schema.optional(_FilePartSource), +}) + .annotate({ identifier: "FilePartInput" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type FilePartInput = Types.DeepMutable> + +export const AgentPartInput = Schema.Struct({ + id: Schema.optional(PartID), + type: Schema.Literal("agent"), + name: Schema.String, + source: Schema.optional( + Schema.Struct({ + value: Schema.String, + start: Schema.Number.check(Schema.isInt()), + end: Schema.Number.check(Schema.isInt()), + }), + ), +}) + .annotate({ identifier: "AgentPartInput" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type AgentPartInput = Types.DeepMutable> + +export const SubtaskPartInput = Schema.Struct({ + id: Schema.optional(PartID), + type: Schema.Literal("subtask"), + prompt: Schema.String, + description: Schema.String, + agent: Schema.String, + model: Schema.optional( + Schema.Struct({ + providerID: ProviderID, + modelID: ModelID, + }), + ), + command: Schema.optional(Schema.String), +}) + .annotate({ identifier: "SubtaskPartInput" }) + .pipe(withStatics((s) => ({ zod: zod(s) }))) +export type SubtaskPartInput = Types.DeepMutable> export const Assistant = Base.extend({ role: z.literal("assistant"), @@ -517,7 +617,10 @@ export const WithParts = z.object({ info: Info, parts: z.array(Part), }) -export type WithParts = z.infer +export type WithParts = { + info: Info + parts: Part[] +} const Cursor = z.object({ id: MessageID.zod, diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index 6dcec04592..9d50db4afb 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -1721,50 +1721,25 @@ export const PromptInput = z.object({ variant: z.string().optional(), parts: z.array( z.discriminatedUnion("type", [ - MessageV2.TextPart.omit({ - messageID: true, - sessionID: true, - }) - .partial({ - id: true, - }) - .meta({ - ref: "TextPartInput", - }), - MessageV2.FilePart.omit({ - messageID: true, - sessionID: true, - }) - .partial({ - id: true, - }) - .meta({ - ref: "FilePartInput", - }), - MessageV2.AgentPart.omit({ - messageID: true, - sessionID: true, - }) - .partial({ - id: true, - }) - .meta({ - ref: "AgentPartInput", - }), - MessageV2.SubtaskPart.omit({ - messageID: true, - sessionID: true, - }) - .partial({ - id: true, - }) - .meta({ - ref: "SubtaskPartInput", - }), + MessageV2.TextPartInput.zod as unknown as z.ZodObject, + MessageV2.FilePartInput.zod as unknown as z.ZodObject, + MessageV2.AgentPartInput.zod as unknown as z.ZodObject, + MessageV2.SubtaskPartInput.zod as unknown as z.ZodObject, ]), ), }) -export type PromptInput = z.infer +// `z.discriminatedUnion` erases the discriminated members' shapes back to +// `{}` because the derived `.zod` on each input is typed as an opaque +// `z.ZodType`. Restore the precise `parts` type from the exported Schema +// input types so callers see a proper tagged union. +type PartInputUnion = + | MessageV2.TextPartInput + | MessageV2.FilePartInput + | MessageV2.AgentPartInput + | MessageV2.SubtaskPartInput +export type PromptInput = Omit, "parts"> & { + parts: PartInputUnion[] +} export const LoopInput = z.object({ sessionID: SessionID.zod, @@ -1792,14 +1767,19 @@ export const CommandInput = z.object({ arguments: z.string(), command: z.string(), variant: z.string().optional(), + // Inlined (no `.meta({ ref })`) to keep the original SDK output — the + // PromptInput call site below references FilePartInput by ref via the + // Schema export in message-v2.ts. parts: z .array( z.discriminatedUnion("type", [ - MessageV2.FilePart.omit({ - messageID: true, - sessionID: true, - }).partial({ - id: true, + z.object({ + id: PartID.zod.optional(), + type: z.literal("file"), + mime: z.string(), + filename: z.string().optional(), + url: z.string(), + source: MessageV2.FilePartSource.zod.optional(), }), ]), )