generate message type

2 files changed, 150 insertions, 3 deletions

diff --git a/js/src/server/message.ts b/js/src/server/message.ts
new file mode 100644
index 000000000..77ef5b8f2
--- /dev/null
+++ b/js/src/server/message.ts
@@ -0,0 +1,146 @@
+import z from "zod";
+
+const ToolCall = z
+  .object({
+    state: z.literal("call"),
+    step: z.number().optional(),
+    toolCallId: z.string(),
+    toolName: z.string(),
+    args: z.record(z.string(), z.any()),
+  })
+  .openapi({
+    ref: "Session.Message.ToolInvocation.ToolCall",
+  });
+
+const ToolPartialCall = z
+  .object({
+    state: z.literal("partial-call"),
+    step: z.number().optional(),
+    toolCallId: z.string(),
+    toolName: z.string(),
+    args: z.record(z.string(), z.any()),
+  })
+  .openapi({
+    ref: "Session.Message.ToolInvocation.ToolPartialCall",
+  });
+
+const ToolResult = z
+  .object({
+    state: z.literal("result"),
+    step: z.number().optional(),
+    toolCallId: z.string(),
+    toolName: z.string(),
+    args: z.record(z.string(), z.any()),
+    result: z.string(),
+  })
+  .openapi({
+    ref: "Session.Message.ToolInvocation.ToolResult",
+  });
+
+const ToolInvocation = z
+  .union([ToolCall, ToolPartialCall, ToolResult])
+  .openapi({
+    ref: "Session.Message.ToolInvocation",
+  });
+export type ToolInvocation = z.infer<typeof ToolInvocation>;
+
+const TextPart = z
+  .object({
+    type: z.literal("text"),
+    text: z.string(),
+  })
+  .openapi({
+    ref: "Session.Message.Part.Text",
+  });
+
+const ReasoningPart = z
+  .object({
+    type: z.literal("reasoning"),
+    text: z.string(),
+    providerMetadata: z.record(z.any()).optional(),
+  })
+  .openapi({
+    ref: "Session.Message.Part.Reasoning",
+  });
+
+const ToolInvocationPart = z
+  .object({
+    type: z.literal("tool-invocation"),
+    toolInvocation: ToolInvocation,
+  })
+  .openapi({
+    ref: "Session.Message.Part.ToolInvocation",
+  });
+
+const SourceUrlPart = z
+  .object({
+    type: z.literal("source-url"),
+    sourceId: z.string(),
+    url: z.string(),
+    title: z.string().optional(),
+    providerMetadata: z.record(z.any()).optional(),
+  })
+  .openapi({
+    ref: "Session.Message.Part.SourceUrl",
+  });
+
+const FilePart = z
+  .object({
+    type: z.literal("file"),
+    mediaType: z.string(),
+    filename: z.string().optional(),
+    url: z.string(),
+  })
+  .openapi({
+    ref: "Session.Message.Part.File",
+  });
+
+const StepStartPart = z
+  .object({
+    type: z.literal("step-start"),
+  })
+  .openapi({
+    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([
+    TextPart,
+    ReasoningPart,
+    ToolInvocationPart,
+    SourceUrlPart,
+    FilePart,
+    StepStartPart,
+    DataPart,
+  ])
+  .openapi({
+    ref: "Session.Message.Part",
+  });
+
+export const SessionMessage = z
+  .object({
+    id: z.string(),
+    role: z.enum(["system", "user", "assistant"]),
+    parts: z.array(Part),
+    metadata: z.object({
+      time: z.object({
+        created: z.number(),
+        completed: z.number().optional(),
+      }),
+      sessionID: z.string(),
+      tool: z.record(z.string(), z.any()),
+    }),
+  })
+  .openapi({
+    ref: "Session.Message",
+  });
diff --git a/js/src/server/server.ts b/js/src/server/server.ts
index 57f52a37c..ea70cd5ef 100644
--- a/js/src/server/server.ts
+++ b/js/src/server/server.ts
@@ -9,6 +9,7 @@ import { z } from "zod";
 import "zod-openapi/extend";
 import { Config } from "../app/config";
 import { LLM } from "../llm/llm";
+import { SessionMessage } from "./message";
 
 const SessionInfo = Session.Info.openapi({
   ref: "Session.Info",
@@ -118,7 +119,7 @@ export namespace Server {
               description: "Successfully created session",
               content: {
                 "application/json": {
-                  schema: resolver(z.custom<Session.Message[]>()),
+                  schema: resolver(SessionMessage.array()),
                 },
               },
             },
@@ -165,12 +166,12 @@ export namespace Server {
             sessionID: z.string(),
             providerID: z.string(),
             modelID: z.string(),
-            parts: z.custom<Session.Message["parts"]>(),
+            parts: SessionMessage.shape.parts,
           }),
         ),
         async (c) => {
           const body = c.req.valid("json");
-          const msg = await Session.chat(body);
+          const msg = await Session.chat(body as any);
           return c.json(msg);
         },
       )