rework types

1 files changed, 160 insertions, 0 deletions

diff --git a/js/src/session/message.ts b/js/src/session/message.ts
new file mode 100644
index 000000000..75c22ef0b
--- /dev/null
+++ b/js/src/session/message.ts
@@ -0,0 +1,160 @@
+import z from "zod";
+
+export namespace Message {
+  export const ToolCall = z
+    .object({
+      state: z.literal("call"),
+      step: z.number().optional(),
+      toolCallId: z.string(),
+      toolName: z.string(),
+      args: z.custom<Required<unknown>>(),
+    })
+    .openapi({
+      ref: "Message.ToolInvocation.ToolCall",
+    });
+  export type ToolCall = z.infer<typeof ToolCall>;
+
+  export const ToolPartialCall = z
+    .object({
+      state: z.literal("partial-call"),
+      step: z.number().optional(),
+      toolCallId: z.string(),
+      toolName: z.string(),
+      args: z.custom<Required<unknown>>(),
+    })
+    .openapi({
+      ref: "Message.ToolInvocation.ToolPartialCall",
+    });
+  export type ToolPartialCall = z.infer<typeof ToolPartialCall>;
+
+  export const ToolResult = z
+    .object({
+      state: z.literal("result"),
+      step: z.number().optional(),
+      toolCallId: z.string(),
+      toolName: z.string(),
+      args: z.custom<Required<unknown>>(),
+      result: z.string(),
+    })
+    .openapi({
+      ref: "Message.ToolInvocation.ToolResult",
+    });
+  export type ToolResult = z.infer<typeof ToolResult>;
+
+  export const ToolInvocation = z
+    .discriminatedUnion("state", [ToolCall, ToolPartialCall, ToolResult])
+    .openapi({
+      ref: "Message.ToolInvocation",
+    });
+  export type ToolInvocation = z.infer<typeof ToolInvocation>;
+
+  export const TextPart = z
+    .object({
+      type: z.literal("text"),
+      text: z.string(),
+    })
+    .openapi({
+      ref: "Message.Part.Text",
+    });
+  export type TextPart = z.infer<typeof TextPart>;
+
+  export const ReasoningPart = z
+    .object({
+      type: z.literal("reasoning"),
+      text: z.string(),
+      providerMetadata: z.record(z.any()).optional(),
+    })
+    .openapi({
+      ref: "Message.Part.Reasoning",
+    });
+  export type ReasoningPart = z.infer<typeof ReasoningPart>;
+
+  export const ToolInvocationPart = z
+    .object({
+      type: z.literal("tool-invocation"),
+      toolInvocation: ToolInvocation,
+    })
+    .openapi({
+      ref: "Message.Part.ToolInvocation",
+    });
+  export type ToolInvocationPart = z.infer<typeof ToolInvocationPart>;
+
+  export 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: "Message.Part.SourceUrl",
+    });
+  export type SourceUrlPart = z.infer<typeof SourceUrlPart>;
+
+  export const FilePart = z
+    .object({
+      type: z.literal("file"),
+      mediaType: z.string(),
+      filename: z.string().optional(),
+      url: z.string(),
+    })
+    .openapi({
+      ref: "Message.Part.File",
+    });
+  export type FilePart = z.infer<typeof FilePart>;
+
+  export const StepStartPart = z
+    .object({
+      type: z.literal("step-start"),
+    })
+    .openapi({
+      ref: "Message.Part.StepStart",
+    });
+  export type StepStartPart = z.infer<typeof StepStartPart>;
+
+  export const Part = z
+    .discriminatedUnion("type", [
+      TextPart,
+      ReasoningPart,
+      ToolInvocationPart,
+      SourceUrlPart,
+      FilePart,
+      StepStartPart,
+    ])
+    .openapi({
+      ref: "Message.Part",
+    });
+  export type Part = z.infer<typeof Part>;
+
+  export const Info = 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()),
+        assistant: z
+          .object({
+            modelID: z.string(),
+            providerID: z.string(),
+            cost: z.number(),
+            tokens: z.object({
+              input: z.number(),
+              output: z.number(),
+              reasoning: z.number(),
+            }),
+          })
+          .optional(),
+      }),
+    })
+    .openapi({
+      ref: "Message.Info",
+    });
+  export type Info = z.infer<typeof Info>;
+}