Merge pull request #213 from get-convex/seth/userid-in-uimessage-188

Add userId field to UIMessage type

Author: sethconvex

src/UIMessages.ts

@@ -42,6 +42,7 @@ export type UIMessage<
   stepOrder: number;
   status: UIStatus;
   agentName?: string;
+  userId?: string;
   text: string;
   _creationTime: number;
 };
@@ -75,7 +76,8 @@ export function fromUIMessages<METADATA = unknown>(
         "providerOptions",
         "metadata",
       ]),
-      ...omit(uiMessage, ["parts", "role", "key", "text"]),
+      ...omit(uiMessage, ["parts", "role", "key", "text", "userId"]),
+      userId: uiMessage.userId ?? meta.userId,
       status: uiMessage.status === "streaming" ? "pending" : "success",
       streaming: uiMessage.status === "streaming",
       // to override
@@ -288,6 +290,7 @@ function createSystemUIMessage<
     text,
     role: "system",
     agentName: message.agentName,
+    userId: message.userId,
     parts: [{ type: "text", text, ...partCommon } satisfies TextUIPart],
     metadata: message.metadata,
   };
@@ -347,6 +350,7 @@ function createUserUIMessage<
     key: `${message.threadId}-${message.order}-${message.stepOrder}`,
     text,
     role: "user",
+    userId: message.userId,
     parts,
     metadata: message.metadata,
   };
@@ -370,6 +374,7 @@ function createAssistantUIMessage<
     stepOrder: firstMessage.stepOrder,
     key: `${firstMessage.threadId}-${firstMessage.order}-${firstMessage.stepOrder}`,
     agentName: firstMessage.agentName,
+    userId: firstMessage.userId,
   };
 
   // Get status from last message

src/toUIMessages.test.ts

@@ -728,4 +728,131 @@ describe("toUIMessages", () => {
     expect(textParts).toHaveLength(1);
     expect(textParts[0].text).toBe("The result is 42.");
   });
+
+  describe("userId preservation", () => {
+    it("preserves userId in user messages", () => {
+      const messages = [
+        baseMessageDoc({
+          userId: "user123",
+          message: {
+            role: "user",
+            content: "Hello!",
+          },
+          text: "Hello!",
+        }),
+      ];
+      const uiMessages = toUIMessages(messages);
+      expect(uiMessages).toHaveLength(1);
+      expect(uiMessages[0].role).toBe("user");
+      expect(uiMessages[0].userId).toBe("user123");
+    });
+
+    it("preserves userId in system messages", () => {
+      const messages = [
+        baseMessageDoc({
+          userId: "user456",
+          message: {
+            role: "system",
+            content: "System prompt",
+          },
+          text: "System prompt",
+        }),
+      ];
+      const uiMessages = toUIMessages(messages);
+      expect(uiMessages).toHaveLength(1);
+      expect(uiMessages[0].role).toBe("system");
+      expect(uiMessages[0].userId).toBe("user456");
+    });
+
+    it("preserves userId in assistant messages", () => {
+      const messages = [
+        baseMessageDoc({
+          userId: "user789",
+          message: {
+            role: "assistant",
+            content: "Hi there!",
+          },
+          text: "Hi there!",
+        }),
+      ];
+      const uiMessages = toUIMessages(messages);
+      expect(uiMessages).toHaveLength(1);
+      expect(uiMessages[0].role).toBe("assistant");
+      expect(uiMessages[0].userId).toBe("user789");
+    });
+
+    it("preserves userId from first message in grouped assistant messages", () => {
+      const messages = [
+        baseMessageDoc({
+          _id: "msg1",
+          userId: "userA",
+          order: 1,
+          stepOrder: 1,
+          tool: true,
+          message: {
+            role: "assistant",
+            content: [
+              {
+                type: "tool-call",
+                toolName: "myTool",
+                toolCallId: "call1",
+                args: {},
+              },
+            ],
+          },
+          text: "",
+        }),
+        baseMessageDoc({
+          _id: "msg2",
+          userId: "userA",
+          order: 1,
+          stepOrder: 2,
+          tool: true,
+          message: {
+            role: "tool",
+            content: [
+              {
+                type: "tool-result",
+                toolCallId: "call1",
+                toolName: "myTool",
+                output: { type: "text", value: "result" },
+              },
+            ],
+          },
+          text: "",
+        }),
+        baseMessageDoc({
+          _id: "msg3",
+          userId: "userA",
+          order: 1,
+          stepOrder: 3,
+          message: {
+            role: "assistant",
+            content: "Done!",
+          },
+          text: "Done!",
+        }),
+      ];
+      const uiMessages = toUIMessages(messages);
+      expect(uiMessages).toHaveLength(1);
+      expect(uiMessages[0].role).toBe("assistant");
+      expect(uiMessages[0].userId).toBe("userA");
+    });
+
+    it("handles undefined userId gracefully", () => {
+      const messages = [
+        baseMessageDoc({
+          // No userId provided
+          message: {
+            role: "user",
+            content: "Hello!",
+          },
+          text: "Hello!",
+        }),
+      ];
+      const uiMessages = toUIMessages(messages);
+      expect(uiMessages).toHaveLength(1);
+      expect(uiMessages[0].userId).toBeUndefined();
+    });
+  });
 });