From 1ae552a92e16df03bb89446d6ae09cfef80fc546 Mon Sep 17 00:00:00 2001
From: theadriann <adrian.brojbeanu@gmail.com>
Date: Fri, 5 Jun 2026 18:45:37 +0700
Subject: [PATCH] feat: allow stream message IDs

---
 src/helpers.ts        |  8 ++++++++
 tests/session.test.ts | 27 +++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/src/helpers.ts b/src/helpers.ts
index f221406..b53d3ca 100644
--- a/src/helpers.ts
+++ b/src/helpers.ts
@@ -66,6 +66,12 @@ export function throwIfAborted(signal: AbortSignal | undefined): void {
 }
 
 export interface MessageOptions {
+  /**
+   * Optional client-generated message id forwarded to add_user_message.
+   * Use this when callers need to correlate later session notifications with
+   * the user message that started this stream.
+   */
+  messageId?: string;
   images?: Base64ImageSource[];
   files?: DocumentSource[];
   outputFormat?: OutputFormat;
@@ -77,6 +83,7 @@ interface StreamableClient {
   onNotification(handler: (n: Record<string, unknown>) => void): () => void;
   addUserMessage(params: {
     text: string;
+    messageId?: string;
     images?: Base64ImageSource[];
     files?: DocumentSource[];
     outputFormat?: OutputFormat;
@@ -120,6 +127,7 @@ export async function* streamFromClient(
     await Promise.race([
       client.addUserMessage({
         text: prompt,
+        messageId: options?.messageId,
         images: options?.images,
         files: options?.files,
         outputFormat: options?.outputFormat,
diff --git a/tests/session.test.ts b/tests/session.test.ts
index 913bf7b..54f8402 100644
--- a/tests/session.test.ts
+++ b/tests/session.test.ts
@@ -443,6 +443,33 @@ describe('DroidSession', () => {
       await session.close();
     });
 
+    it('forwards an explicit stream messageId to add_user_message', async () => {
+      const transport = new InMemoryTransport();
+      await transport.connect();
+
+      setupFullResponder(transport, 'sess-stream-message-id');
+
+      const session = await createSession({ apiKey: 'test-key', transport });
+
+      for await (const msg of session.stream('Hello', {
+        messageId: 'message-explicit-1',
+      })) {
+        void msg;
+      }
+
+      const addMessageRequest = transport.sentMessages.find(
+        (message) =>
+          (message as Record<string, unknown>)['method'] ===
+          DroidServerMethod.ADD_USER_MESSAGE
+      ) as Record<string, unknown>;
+      const params = addMessageRequest['params'] as Record<string, unknown>;
+
+      expect(params['text']).toBe('Hello');
+      expect(params['messageId']).toBe('message-explicit-1');
+
+      await session.close();
+    });
+
     it('defaults to message-level events and opts into partial events', async () => {
       const createStreamingSession = async (
         sessionId: string