From 17426a1504a9eb88f7bb4732227d7981fe253e54 Mon Sep 17 00:00:00 2001 From: User Date: Tue, 19 May 2026 13:30:36 -0700 Subject: [PATCH 1/3] feat: support custom message ids Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- examples/message-id-run.ts | 47 +++++++++++++++++++++++++++++ examples/message-id-stream.ts | 56 +++++++++++++++++++++++++++++++++++ src/session.ts | 4 +++ tests/run.test.ts | 2 ++ tests/session.test.ts | 50 +++++++++++++++++++++++++++++++ 5 files changed, 159 insertions(+) create mode 100644 examples/message-id-run.ts create mode 100644 examples/message-id-stream.ts diff --git a/examples/message-id-run.ts b/examples/message-id-run.ts new file mode 100644 index 0000000..51db3aa --- /dev/null +++ b/examples/message-id-run.ts @@ -0,0 +1,47 @@ +/** + * Manual smoke test for `run(prompt, { messageId })`. + * + * Usage: + * npx tsx examples/message-id-run.ts + * npx tsx examples/message-id-run.ts "Reply with exactly: RUN_OK" + */ + +import { randomUUID } from 'node:crypto'; + +import { DroidMessageType, run } from '../src/index.js'; + +async function main(): Promise { + const prompt = + process.argv.slice(2).join(' ') || 'Reply with exactly: RUN_OK'; + const messageId = `sdk-run-${randomUUID()}`; + + const result = await run(prompt, { + cwd: process.cwd(), + messageId, + }); + const userMessage = result.messages.find( + (msg) => msg.type === DroidMessageType.User + ); + const observedUserMessageId = userMessage?.message.id; + + if (observedUserMessageId !== messageId) { + throw new Error( + `Expected user messageId ${messageId}, got ${observedUserMessageId ?? 'none'}` + ); + } + + console.log( + JSON.stringify({ + api: 'run', + sessionId: result.sessionId, + messageId, + observedUserMessageId, + text: result.text, + }) + ); +} + +main().catch((err: unknown) => { + console.error('Error:', err); + process.exit(1); +}); diff --git a/examples/message-id-stream.ts b/examples/message-id-stream.ts new file mode 100644 index 0000000..f5645e0 --- /dev/null +++ b/examples/message-id-stream.ts @@ -0,0 +1,56 @@ +/** + * Manual smoke test for `session.stream(prompt, { messageId })`. + * + * Usage: + * npx tsx examples/message-id-stream.ts + * npx tsx examples/message-id-stream.ts "Reply with exactly: STREAM_OK" + */ + +import { randomUUID } from 'node:crypto'; + +import { DroidMessageType, createSession } from '../src/index.js'; + +async function main(): Promise { + const prompt = + process.argv.slice(2).join(' ') || 'Reply with exactly: STREAM_OK'; + const messageId = `sdk-stream-${randomUUID()}`; + const session = await createSession({ cwd: process.cwd() }); + + try { + let observedUserMessageId: string | undefined; + let text = ''; + + for await (const msg of session.stream(prompt, { messageId })) { + if (msg.type === DroidMessageType.User) { + observedUserMessageId = msg.message.id; + } else if (msg.type === DroidMessageType.Assistant) { + text += msg.text; + } else if (msg.type === DroidMessageType.Result && text.length === 0) { + text = msg.text; + } + } + + if (observedUserMessageId !== messageId) { + throw new Error( + `Expected user messageId ${messageId}, got ${observedUserMessageId ?? 'none'}` + ); + } + + console.log( + JSON.stringify({ + api: 'session.stream', + sessionId: session.sessionId, + messageId, + observedUserMessageId, + text, + }) + ); + } finally { + await session.close(); + } +} + +main().catch((err: unknown) => { + console.error('Error:', err); + process.exit(1); +}); diff --git a/src/session.ts b/src/session.ts index 0614055..a25b545 100644 --- a/src/session.ts +++ b/src/session.ts @@ -78,6 +78,7 @@ export interface ResumeSessionOptions extends Pick< } export interface MessageOptions { + messageId?: string; images?: Base64ImageSource[]; files?: DocumentSource[]; outputFormat?: OutputFormat; @@ -209,6 +210,9 @@ export class DroidSession { try { await Promise.race([ this._client.addUserMessage({ + ...(options?.messageId !== undefined && { + messageId: options.messageId, + }), text: prompt, images: options?.images, files: options?.files, diff --git a/tests/run.test.ts b/tests/run.test.ts index 06e1198..a0557c1 100644 --- a/tests/run.test.ts +++ b/tests/run.test.ts @@ -78,6 +78,7 @@ describe('run()', () => { machineId: 'machine-1', modelId: 'model-1', reasoningEffort: ReasoningEffort.High, + messageId: 'run-message-id', images: [{ type: 'base64', data: 'image-data', mediaType: 'image/png' }], files: [ { @@ -106,6 +107,7 @@ describe('run()', () => { DroidServerMethod.ADD_USER_MESSAGE ) as Record; const addParams = addMsg['params'] as Record; + expect(addParams['messageId']).toBe('run-message-id'); expect(addParams['text']).toBe('Describe these inputs'); expect(addParams['images']).toEqual([ { type: 'base64', data: 'image-data', mediaType: 'image/png' }, diff --git a/tests/session.test.ts b/tests/session.test.ts index 64cbca4..81da8b6 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -362,6 +362,56 @@ describe('DroidSession', () => { await session.close(); }); + it('passes custom messageId in addUserMessage RPC params', async () => { + const transport = new InMemoryTransport(); + await transport.connect(); + + setupFullResponder(transport, 'sess-stream-message-id'); + + const session = await createSession({ transport }); + + for await (const _msg of session.stream('Hello', { + messageId: 'caller-message-id', + })) { + void _msg; + } + + const addMsg = transport.sentMessages.find( + (message) => + (message as Record)['method'] === + DroidServerMethod.ADD_USER_MESSAGE + ) as Record; + const addParams = addMsg['params'] as Record; + expect(addParams['messageId']).toBe('caller-message-id'); + expect(addParams['text']).toBe('Hello'); + + await session.close(); + }); + + it('omits messageId from addUserMessage RPC params by default', async () => { + const transport = new InMemoryTransport(); + await transport.connect(); + + setupFullResponder(transport, 'sess-stream-default-message-id'); + + const session = await createSession({ transport }); + + for await (const _msg of session.stream('Hello')) { + void _msg; + } + + const addMsg = transport.sentMessages.find( + (message) => + (message as Record)['method'] === + DroidServerMethod.ADD_USER_MESSAGE + ) as Record; + const addParams = addMsg['params'] as Record; + expect(addParams).not.toHaveProperty('messageId'); + expect(addParams['text']).toBe('Hello'); + + await session.close(); + }); + it('defaults to message-level events and opts into partial events', async () => { const createStreamingSession = async ( sessionId: string From c8773eb2c6310a4a599c6a92210816fa4f4755a5 Mon Sep 17 00:00:00 2001 From: User Date: Tue, 19 May 2026 13:52:29 -0700 Subject: [PATCH 2/3] fix: validate custom message ids Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/client.ts | 3 ++- src/schemas/client.ts | 9 ++++++++- tests/session.test.ts | 34 ++++++++++++++++++++++++++++++++++ 3 files changed, 44 insertions(+), 2 deletions(-) diff --git a/src/client.ts b/src/client.ts index 42fdf05..49f994d 100644 --- a/src/client.ts +++ b/src/client.ts @@ -59,6 +59,7 @@ import type { } from './schemas/client.js'; import { AddMcpServerResultSchema, + AddUserMessageRequestParamsSchema, AddUserMessageResultSchema, AuthenticateMcpServerResultSchema, CancelMcpAuthResultSchema, @@ -232,7 +233,7 @@ export class DroidClient { ): Promise { return this._sessionRpc( DroidServerMethod.ADD_USER_MESSAGE, - params, + AddUserMessageRequestParamsSchema.parse(params), AddUserMessageResultSchema ); } diff --git a/src/schemas/client.ts b/src/schemas/client.ts index 00c58c6..8269c45 100644 --- a/src/schemas/client.ts +++ b/src/schemas/client.ts @@ -289,10 +289,17 @@ export const OutputFormatSchema = z export type OutputFormat = z.infer; +export const MessageIdSchema = z + .string() + .max(512, 'messageId must be at most 512 characters') + .refine((value) => value.trim().length > 0, { + message: 'messageId must be a non-empty string', + }); + /** Parameters for droid.add_user_message request. */ export const AddUserMessageRequestParamsSchema = z .object({ - messageId: z.string().optional(), + messageId: MessageIdSchema.optional(), text: z.string(), images: z.array(Base64ImageSourceSchema).optional(), files: z.array(DocumentSourceSchema).optional(), diff --git a/tests/session.test.ts b/tests/session.test.ts index 81da8b6..f2d0973 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -412,6 +412,40 @@ describe('DroidSession', () => { await session.close(); }); + it.each([ + ['empty string', ''], + ['whitespace-only string', ' '], + ['too-long string', 'x'.repeat(513)], + ['non-string value', 123], + ])('rejects invalid messageId: %s', async (_label, messageId) => { + const transport = new InMemoryTransport(); + await transport.connect(); + + setupFullResponder(transport, 'sess-stream-invalid-message-id'); + + const session = await createSession({ transport }); + + await expect( + (async () => { + for await (const _msg of session.stream('Hello', { + messageId: messageId as never, + })) { + void _msg; + } + })() + ).rejects.toThrow(); + + expect( + transport.sentMessages.some( + (message) => + (message as Record)['method'] === + DroidServerMethod.ADD_USER_MESSAGE + ) + ).toBe(false); + + await session.close(); + }); + it('defaults to message-level events and opts into partial events', async () => { const createStreamingSession = async ( sessionId: string From 29e562be37b7cc67f4a152ce5d9ae07033c0512d Mon Sep 17 00:00:00 2001 From: User Date: Tue, 19 May 2026 16:51:36 -0700 Subject: [PATCH 3/3] fix: simplify message id validation Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- examples/fork-session.ts | 2 +- examples/interrupt-session.ts | 2 +- examples/multi-turn-session.ts | 2 +- examples/sdk-mcp-tool.ts | 2 +- src/client.ts | 8 ++++++-- src/schemas/client.ts | 4 ++-- tests/helpers.ts | 19 +++++++++++++++++++ tests/run.test.ts | 11 +++++------ tests/session.test.ts | 23 ++++++++++------------- 9 files changed, 46 insertions(+), 27 deletions(-) diff --git a/examples/fork-session.ts b/examples/fork-session.ts index c35524f..2a17fbd 100644 --- a/examples/fork-session.ts +++ b/examples/fork-session.ts @@ -22,7 +22,7 @@ async function streamText( ): Promise { let text = ''; for await (const msg of session.stream(prompt)) { - if (msg.type === DroidMessageType.AssistantTextDelta) { + if (msg.type === DroidMessageType.Assistant) { text += msg.text; } } diff --git a/examples/interrupt-session.ts b/examples/interrupt-session.ts index 02d46e3..fc8be74 100644 --- a/examples/interrupt-session.ts +++ b/examples/interrupt-session.ts @@ -14,7 +14,7 @@ try { for await (const msg of session.stream( 'Write a long history of computing.' )) { - if (msg.type !== DroidMessageType.AssistantTextDelta) { + if (msg.type !== DroidMessageType.Assistant) { continue; } diff --git a/examples/multi-turn-session.ts b/examples/multi-turn-session.ts index 7cf3feb..c96d11b 100644 --- a/examples/multi-turn-session.ts +++ b/examples/multi-turn-session.ts @@ -13,7 +13,7 @@ async function streamText( ): Promise { let text = ''; for await (const msg of session.stream(prompt)) { - if (msg.type === DroidMessageType.AssistantTextDelta) { + if (msg.type === DroidMessageType.Assistant) { text += msg.text; } } diff --git a/examples/sdk-mcp-tool.ts b/examples/sdk-mcp-tool.ts index 66d1bf3..b3c3a75 100644 --- a/examples/sdk-mcp-tool.ts +++ b/examples/sdk-mcp-tool.ts @@ -32,7 +32,7 @@ try { for await (const msg of session.stream( 'Use the favorite_number tool for Ada and tell me the answer.' )) { - if (msg.type === DroidMessageType.AssistantTextDelta) { + if (msg.type === DroidMessageType.Assistant) { process.stdout.write(msg.text); } } diff --git a/src/client.ts b/src/client.ts index 49f994d..9d90d3d 100644 --- a/src/client.ts +++ b/src/client.ts @@ -59,7 +59,6 @@ import type { } from './schemas/client.js'; import { AddMcpServerResultSchema, - AddUserMessageRequestParamsSchema, AddUserMessageResultSchema, AuthenticateMcpServerResultSchema, CancelMcpAuthResultSchema, @@ -80,6 +79,7 @@ import { ListToolsResultSchema, ListSkillsResultSchema, LoadSessionResultSchema, + MessageIdSchema, RemoveMcpServerResultSchema, SubmitBugReportResultSchema, SubmitMcpAuthCodeResultSchema, @@ -231,9 +231,13 @@ export class DroidClient { > > ): Promise { + if (params.messageId !== undefined) { + MessageIdSchema.parse(params.messageId); + } + return this._sessionRpc( DroidServerMethod.ADD_USER_MESSAGE, - AddUserMessageRequestParamsSchema.parse(params), + params, AddUserMessageResultSchema ); } diff --git a/src/schemas/client.ts b/src/schemas/client.ts index 8269c45..03152d3 100644 --- a/src/schemas/client.ts +++ b/src/schemas/client.ts @@ -564,7 +564,7 @@ export type RewindEvictedFile = z.infer; /** Parameters for droid.get_rewind_info request. */ export const GetRewindInfoRequestParamsSchema = z .object({ - messageId: z.string(), + messageId: MessageIdSchema, }) .passthrough(); @@ -575,7 +575,7 @@ export type GetRewindInfoRequestParams = z.infer< /** Parameters for droid.execute_rewind request. */ export const ExecuteRewindRequestParamsSchema = z .object({ - messageId: z.string(), + messageId: MessageIdSchema, filesToRestore: z.array(RewindFileSnapshotSchema), filesToDelete: z.array(RewindFileCreationSchema), forkTitle: z.string(), diff --git a/tests/helpers.ts b/tests/helpers.ts index 3f04686..1c1e334 100644 --- a/tests/helpers.ts +++ b/tests/helpers.ts @@ -114,6 +114,25 @@ export class InMemoryTransport implements DroidClientTransport { } } +export function findSentRequestParams( + transport: InMemoryTransport, + method: string +): Record { + const message = transport.sentMessages.find((sentMessage) => { + return sentMessage['method'] === method; + }); + if (!message) { + throw new Error(`Expected request for method ${method}`); + } + + const params = message['params']; + if (!params || typeof params !== 'object' || Array.isArray(params)) { + throw new Error(`Expected params for method ${method}`); + } + + return params as Record; +} + export function makeSuccessResponse( id: string, result: JsonRpcTestMessage = {} diff --git a/tests/run.test.ts b/tests/run.test.ts index a0557c1..3246da5 100644 --- a/tests/run.test.ts +++ b/tests/run.test.ts @@ -11,6 +11,7 @@ import { } from '../src/schemas/index.js'; import { InMemoryTransport, + findSentRequestParams, makeErrorResponse, makeSessionNotification, makeSuccessResponse, @@ -101,12 +102,10 @@ describe('run()', () => { expect(initParams['modelId']).toBe('model-1'); expect(initParams['reasoningEffort']).toBe(ReasoningEffort.High); - const addMsg = transport.sentMessages.find( - (message) => - (message as Record)['method'] === - DroidServerMethod.ADD_USER_MESSAGE - ) as Record; - const addParams = addMsg['params'] as Record; + const addParams = findSentRequestParams( + transport, + DroidServerMethod.ADD_USER_MESSAGE + ); expect(addParams['messageId']).toBe('run-message-id'); expect(addParams['text']).toBe('Describe these inputs'); expect(addParams['images']).toEqual([ diff --git a/tests/session.test.ts b/tests/session.test.ts index f2d0973..f025cb5 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -27,6 +27,7 @@ import { InMemoryTransport, collectStreamText, findLastResult, + findSentRequestParams, makeErrorResponse, makeSessionNotification, makeSuccessResponse, @@ -376,12 +377,10 @@ describe('DroidSession', () => { void _msg; } - const addMsg = transport.sentMessages.find( - (message) => - (message as Record)['method'] === - DroidServerMethod.ADD_USER_MESSAGE - ) as Record; - const addParams = addMsg['params'] as Record; + const addParams = findSentRequestParams( + transport, + DroidServerMethod.ADD_USER_MESSAGE + ); expect(addParams['messageId']).toBe('caller-message-id'); expect(addParams['text']).toBe('Hello'); @@ -400,12 +399,10 @@ describe('DroidSession', () => { void _msg; } - const addMsg = transport.sentMessages.find( - (message) => - (message as Record)['method'] === - DroidServerMethod.ADD_USER_MESSAGE - ) as Record; - const addParams = addMsg['params'] as Record; + const addParams = findSentRequestParams( + transport, + DroidServerMethod.ADD_USER_MESSAGE + ); expect(addParams).not.toHaveProperty('messageId'); expect(addParams['text']).toBe('Hello'); @@ -428,7 +425,7 @@ describe('DroidSession', () => { await expect( (async () => { for await (const _msg of session.stream('Hello', { - messageId: messageId as never, + messageId: messageId as string, })) { void _msg; }