add function to verify zod schema matches type

app/src/pages/playground/playgroundUtils.ts

@@ -4,7 +4,7 @@ import {
   ChatMessageRole,
   createPlaygroundInstance,
   generateMessageId,
-} from "@phoenix/store/playgroundStore";
+} from "@phoenix/store";
 import { safelyParseJSON } from "@phoenix/utils/jsonUtils";
 
 import { ChatRoleMap, DEFAULT_CHAT_ROLE } from "./constants";

app/src/pages/playground/schemas.ts

@@ -6,7 +6,8 @@ import {
   SemanticAttributePrefixes,
 } from "@arizeai/openinference-semantic-conventions";
 
-import { ChatMessageRole } from "@phoenix/store";
+import { ChatMessage, ChatMessageRole } from "@phoenix/store";
+import { schemaForType } from "@phoenix/typeUtils";
 
 /**
  * The zod schema for llm tool calls in an input message
@@ -77,12 +78,15 @@ export const outputSchema = z.object({
  */
 export const chatMessageRolesSchema = z.nativeEnum(ChatMessageRole);
 
-/**
- * The zod schema for ChatMessages
- */
-export const chatMessagesSchema = z.array(
+const chatMessageSchema = schemaForType<ChatMessage>()(
   z.object({
+    id: z.number(),
     role: chatMessageRolesSchema,
     content: z.string(),
   })
 );
+
+/**
+ * The zod schema for ChatMessages
+ */
+export const chatMessagesSchema = z.array(chatMessageSchema);

app/src/typeUtils.ts

@@ -1,3 +1,5 @@
+import { z } from "zod";
+
 /**
  * Utility function that uses the type system to check if a switch statement is exhaustive.
  * If the switch statement is not exhaustive, there will be a type error caught in typescript
@@ -45,3 +47,25 @@ export function isObject(value: unknown): value is object {
 export type Mutable<T> = {
   -readonly [P in keyof T]: T[P];
 };
+
+/**
+ * A zod type utility that ensures that the schema is written to correctly match (at least) what is included in the type.
+ * Note it does not guard against extra fields in the schema not present in the type.
+ * @example
+ * ```typescript
+ * const chatMessageSchema = schemaForType<ChatMessage>()(
+ *  z.object({
+ *    id: z.number(),
+ *    role: chatMessageRolesSchema,
+ *    content: z.string(),
+ *  })
+ * );
+ * ```
+ * Taken from the zod maintainer here:
+ * @see https://github.com/colinhacks/zod/issues/372#issuecomment-826380330
+ */
+export const schemaForType =
+  <T>() =>
+  <S extends z.ZodType<T>>(arg: S) => {
+    return arg;
+  };