Browse Source

Validate prompt messages with Effect Schema (#26796)

Kit Langton 4 months ago
parent
commit
274033cd52

+ 7 - 26
packages/opencode/src/session/message-v2.ts

@@ -23,7 +23,7 @@ import type { SystemError } from "bun"
 import type { Provider } from "@/provider/provider"
 import type { Provider } from "@/provider/provider"
 import { ModelID, ProviderID } from "@/provider/schema"
 import { ModelID, ProviderID } from "@/provider/schema"
 import { Effect, Schema, Types } from "effect"
 import { Effect, Schema, Types } from "effect"
-import { zod, ZodOverride } from "@opencode-ai/core/effect-zod"
+import { zod } from "@opencode-ai/core/effect-zod"
 import { NonNegativeInt, withStatics } from "@opencode-ai/core/schema"
 import { NonNegativeInt, withStatics } from "@opencode-ai/core/schema"
 import { namedSchemaError } from "@/util/named-schema-error"
 import { namedSchemaError } from "@/util/named-schema-error"
 import * as EffectLogger from "@opencode-ai/core/effect/logger"
 import * as EffectLogger from "@opencode-ai/core/effect/logger"
@@ -402,7 +402,7 @@ export const User = Schema.Struct({
   .pipe(withStatics((s) => ({ zod: zod(s) })))
   .pipe(withStatics((s) => ({ zod: zod(s) })))
 export type User = Types.DeepMutable<Schema.Schema.Type<typeof User>>
 export type User = Types.DeepMutable<Schema.Schema.Type<typeof User>>
 
 
-const _Part = Schema.Union([
+export const Part = Schema.Union([
   TextPart,
   TextPart,
   SubtaskPart,
   SubtaskPart,
   ReasoningPart,
   ReasoningPart,
@@ -416,22 +416,6 @@ const _Part = Schema.Union([
   RetryPart,
   RetryPart,
   CompactionPart,
   CompactionPart,
 ]).annotate({ discriminator: "type", identifier: "Part" })
 ]).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
@@ -573,15 +557,12 @@ export type Assistant = Omit<Types.DeepMutable<Schema.Schema.Type<typeof Assista
   error?: AssistantError
   error?: AssistantError
 }
 }
 
 
-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<User | Assistant>,
-})
+export const Info = Schema.Union([User, Assistant]).annotate({ discriminator: "role", identifier: "Message" })
 export type Info = User | Assistant
 export type Info = User | Assistant
 
 
 const UpdatedEventSchema = Schema.Struct({
 const UpdatedEventSchema = Schema.Struct({
   sessionID: SessionID,
   sessionID: SessionID,
-  info: _Info,
+  info: Info,
 })
 })
 
 
 const RemovedEventSchema = Schema.Struct({
 const RemovedEventSchema = Schema.Struct({
@@ -591,7 +572,7 @@ const RemovedEventSchema = Schema.Struct({
 
 
 const PartUpdatedEventSchema = Schema.Struct({
 const PartUpdatedEventSchema = Schema.Struct({
   sessionID: SessionID,
   sessionID: SessionID,
-  part: _Part,
+  part: Part,
   time: NonNegativeInt,
   time: NonNegativeInt,
 })
 })
 
 
@@ -639,8 +620,8 @@ export const Event = {
 }
 }
 
 
 export const WithParts = Schema.Struct({
 export const WithParts = Schema.Struct({
-  info: _Info,
-  parts: Schema.Array(_Part),
+  info: Info,
+  parts: Schema.Array(Part),
 }).pipe(withStatics((s) => ({ zod: zod(s) })))
 }).pipe(withStatics((s) => ({ zod: zod(s) })))
 export type WithParts = {
 export type WithParts = {
   info: Info
   info: Info

+ 9 - 6
packages/opencode/src/session/prompt.ts

@@ -65,6 +65,9 @@ import { SessionTable } from "./session.sql"
 // @ts-ignore
 // @ts-ignore
 globalThis.AI_SDK_LOG_WARNINGS = false
 globalThis.AI_SDK_LOG_WARNINGS = false
 
 
+const decodeMessageInfo = Schema.decodeUnknownExit(MessageV2.Info)
+const decodeMessagePart = Schema.decodeUnknownExit(MessageV2.Part)
+
 const STRUCTURED_OUTPUT_DESCRIPTION = `Use this tool to return your final response in the requested structured format.
 const STRUCTURED_OUTPUT_DESCRIPTION = `Use this tool to return your final response in the requested structured format.
 
 
 IMPORTANT:
 IMPORTANT:
@@ -1292,26 +1295,26 @@ NOTE: At any point in time through this workflow you should feel free to ask the
 
 
       const parts = resolvedParts
       const parts = resolvedParts
 
 
-      const parsed = MessageV2.Info.zod.safeParse(info)
-      if (!parsed.success) {
+      const parsed = decodeMessageInfo(info, { errors: "all", propertyOrder: "original" })
+      if (Exit.isFailure(parsed)) {
         log.error("invalid user message before save", {
         log.error("invalid user message before save", {
           sessionID: input.sessionID,
           sessionID: input.sessionID,
           messageID: info.id,
           messageID: info.id,
           agent: info.agent,
           agent: info.agent,
           model: info.model,
           model: info.model,
-          issues: parsed.error.issues,
+          cause: Cause.pretty(parsed.cause),
         })
         })
       }
       }
       parts.forEach((part, index) => {
       parts.forEach((part, index) => {
-        const p = MessageV2.Part.zod.safeParse(part)
-        if (p.success) return
+        const p = decodeMessagePart(part, { errors: "all", propertyOrder: "original" })
+        if (Exit.isSuccess(p)) return
         log.error("invalid user part before save", {
         log.error("invalid user part before save", {
           sessionID: input.sessionID,
           sessionID: input.sessionID,
           messageID: info.id,
           messageID: info.id,
           partID: part.id,
           partID: part.id,
           partType: part.type,
           partType: part.type,
           index,
           index,
-          issues: p.error.issues,
+          cause: Cause.pretty(p.cause),
           part,
           part,
         })
         })
       })
       })