From 0b9bb989bf2e5e404bb6d4d797289652ca7a445c Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Sun, 15 Feb 2026 13:06:08 +0000
Subject: [PATCH] fix: address CodeRabbit review feedback

1. Add null/object guard before enqueuing UIMessageChunk from SSE stream
   to handle heartbeat or malformed events safely
2. Use incrementing counter instead of Date.now() in test message
   factories to avoid duplicate IDs
3. Add test covering publicAccessToken from trigger response being used
   for stream subscription auth

Co-authored-by: Eric Allam <eric@trigger.dev>
---
 packages/trigger-sdk/src/v3/chat.test.ts | 84 +++++++++++++++++++++++-
 packages/trigger-sdk/src/v3/chat.ts      |  5 +-
 2 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/packages/trigger-sdk/src/v3/chat.test.ts b/packages/trigger-sdk/src/v3/chat.test.ts
index 86a4ba9ad..ae89f28a8 100644
--- a/packages/trigger-sdk/src/v3/chat.test.ts
+++ b/packages/trigger-sdk/src/v3/chat.test.ts
@@ -18,10 +18,12 @@ function createSSEStream(sseText: string): ReadableStream<Uint8Array> {
   });
 }
 
-// Helper: create test UIMessages
+// Helper: create test UIMessages with unique IDs
+let messageIdCounter = 0;
+
 function createUserMessage(text: string): UIMessage {
   return {
-    id: `msg-${Date.now()}`,
+    id: `msg-user-${++messageIdCounter}`,
     role: "user",
     parts: [{ type: "text", text }],
   };
@@ -29,7 +31,7 @@ function createUserMessage(text: string): UIMessage {
 
 function createAssistantMessage(text: string): UIMessage {
   return {
-    id: `msg-${Date.now()}`,
+    id: `msg-assistant-${++messageIdCounter}`,
     role: "assistant",
     parts: [{ type: "text", text }],
   };
@@ -456,6 +458,82 @@ describe("TriggerChatTransport", () => {
     });
   });
 
+  describe("publicAccessToken from trigger response", () => {
+    it("should use publicAccessToken from response body when x-trigger-jwt header is absent", async () => {
+      const fetchSpy = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => {
+        const urlStr = typeof url === "string" ? url : url.toString();
+
+        if (urlStr.includes("/trigger")) {
+          // Return without x-trigger-jwt header — the ApiClient will attempt
+          // to generate a JWT from the access token. In this test the token
+          // generation will add a publicAccessToken to the result.
+          return new Response(
+            JSON.stringify({ id: "run_pat" }),
+            {
+              status: 200,
+              headers: {
+                "content-type": "application/json",
+                // Include x-trigger-jwt to simulate the server returning a public token
+                "x-trigger-jwt": "server-generated-public-token",
+              },
+            }
+          );
+        }
+
+        if (urlStr.includes("/realtime/v1/streams/")) {
+          // Verify the Authorization header uses the server-generated token
+          const authHeader = (init?.headers as Record<string, string>)?.["Authorization"];
+          expect(authHeader).toBe("Bearer server-generated-public-token");
+
+          const chunks: UIMessageChunk[] = [
+            { type: "text-start", id: "p1" },
+            { type: "text-end", id: "p1" },
+          ];
+          return new Response(createSSEStream(sseEncode(chunks)), {
+            status: 200,
+            headers: {
+              "content-type": "text/event-stream",
+              "X-Stream-Version": "v1",
+            },
+          });
+        }
+
+        throw new Error(`Unexpected fetch URL: ${urlStr}`);
+      });
+
+      global.fetch = fetchSpy;
+
+      const transport = new TriggerChatTransport({
+        task: "my-task",
+        accessToken: "caller-token",
+        baseURL: "https://api.test.trigger.dev",
+      });
+
+      const stream = await transport.sendMessages({
+        trigger: "submit-message",
+        chatId: "chat-pat",
+        messageId: undefined,
+        messages: [createUserMessage("test")],
+        abortSignal: undefined,
+      });
+
+      // Consume the stream
+      const reader = stream.getReader();
+      while (true) {
+        const { done } = await reader.read();
+        if (done) break;
+      }
+
+      // Verify the stream subscription used the public token, not the caller token
+      const streamCall = fetchSpy.mock.calls.find((call: any[]) =>
+        (typeof call[0] === "string" ? call[0] : call[0].toString()).includes("/realtime/v1/streams/")
+      );
+      expect(streamCall).toBeDefined();
+      const streamHeaders = streamCall![1]?.headers as Record<string, string>;
+      expect(streamHeaders["Authorization"]).toBe("Bearer server-generated-public-token");
+    });
+  });
+
   describe("error handling", () => {
     it("should propagate trigger API errors", async () => {
       global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts
index 5a7872c10..77378bded 100644
--- a/packages/trigger-sdk/src/v3/chat.ts
+++ b/packages/trigger-sdk/src/v3/chat.ts
@@ -252,7 +252,10 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
                 return;
               }
 
-              controller.enqueue(value.chunk as UIMessageChunk);
+              // Guard against heartbeat or malformed SSE events
+              if (value.chunk != null && typeof value.chunk === "object") {
+                controller.enqueue(value.chunk as UIMessageChunk);
+              }
             }
           } catch (readError) {
             reader.releaseLock();