From f3b224090cf6e719d58d9a36b3e07380aac68d31 Mon Sep 17 00:00:00 2001
From: Dax Raad <d@ironbay.co>
Date: Wed, 28 May 2025 13:35:05 -0400
Subject: better discriminator

---
 js/src/server/message.ts | 15 ++-------------
 1 file changed, 2 insertions(+), 13 deletions(-)

(limited to 'js')

diff --git a/js/src/server/message.ts b/js/src/server/message.ts
index 77ef5b8f2..4ad301a35 100644
--- a/js/src/server/message.ts
+++ b/js/src/server/message.ts
@@ -38,7 +38,7 @@ const ToolResult = z
   });
 
 const ToolInvocation = z
-  .union([ToolCall, ToolPartialCall, ToolResult])
+  .discriminatedUnion("state", [ToolCall, ToolPartialCall, ToolResult])
   .openapi({
     ref: "Session.Message.ToolInvocation",
   });
@@ -103,25 +103,14 @@ const StepStartPart = z
     ref: "Session.Message.Part.StepStart",
   });
 
-const DataPart = z
-  .object({
-    type: z.custom<`data-${string}`>(),
-    id: z.string().optional(),
-    data: z.unknown(),
-  })
-  .openapi({
-    ref: "Session.Message.Part.Data",
-  });
-
 const Part = z
-  .union([
+  .discriminatedUnion("type", [
     TextPart,
     ReasoningPart,
     ToolInvocationPart,
     SourceUrlPart,
     FilePart,
     StepStartPart,
-    DataPart,
   ])
   .openapi({
     ref: "Session.Message.Part",
-- 
cgit v1.2.3