diff --git a/packages/sdk-typescript/src/query/createQuery.ts b/packages/sdk-typescript/src/query/createQuery.ts index e3907635..7549b2b3 100644 --- a/packages/sdk-typescript/src/query/createQuery.ts +++ b/packages/sdk-typescript/src/query/createQuery.ts @@ -43,6 +43,7 @@ export function query({ coreTools: options.coreTools, excludeTools: options.excludeTools, authType: options.authType, + includePartialMessages: options.includePartialMessages, }); const queryOptions: QueryOptions = { diff --git a/packages/sdk-typescript/src/transport/ProcessTransport.ts b/packages/sdk-typescript/src/transport/ProcessTransport.ts index 62a6b2d0..ba13f044 100644 --- a/packages/sdk-typescript/src/transport/ProcessTransport.ts +++ b/packages/sdk-typescript/src/transport/ProcessTransport.ts @@ -155,6 +155,10 @@ export class ProcessTransport implements Transport { args.push('--auth-type', this.options.authType); } + if (this.options.includePartialMessages) { + args.push('--include-partial-messages'); + } + return args; } diff --git a/packages/sdk-typescript/src/types/queryOptionsSchema.ts b/packages/sdk-typescript/src/types/queryOptionsSchema.ts index 7573abef..c347bfdd 100644 --- a/packages/sdk-typescript/src/types/queryOptionsSchema.ts +++ b/packages/sdk-typescript/src/types/queryOptionsSchema.ts @@ -76,6 +76,7 @@ export const QueryOptionsSchema = z ), ) .optional(), + includePartialMessages: z.boolean().optional(), }) .strict(); diff --git a/packages/sdk-typescript/src/types/types.ts b/packages/sdk-typescript/src/types/types.ts index 856099fc..e4cbbb5b 100644 --- a/packages/sdk-typescript/src/types/types.ts +++ b/packages/sdk-typescript/src/types/types.ts @@ -30,6 +30,7 @@ export type TransportOptions = { coreTools?: string[]; excludeTools?: string[]; authType?: string; + includePartialMessages?: boolean; }; type ToolInput = Record; diff --git a/packages/sdk-typescript/test/e2e/multi-turn.test.ts b/packages/sdk-typescript/test/e2e/multi-turn.test.ts index 8e79898e..52c012c8 100644 --- a/packages/sdk-typescript/test/e2e/multi-turn.test.ts +++ b/packages/sdk-typescript/test/e2e/multi-turn.test.ts @@ -476,4 +476,67 @@ describe('Multi-Turn Conversations (E2E)', () => { } }); }); + + describe('Partial Messages in Multi-Turn', () => { + it('should receive partial messages when includePartialMessages is enabled', async () => { + async function* createMultiTurnConversation(): AsyncIterable { + const sessionId = crypto.randomUUID(); + + yield { + type: 'user', + session_id: sessionId, + message: { + role: 'user', + content: 'What is 1 + 1?', + }, + parent_tool_use_id: null, + } as CLIUserMessage; + + await new Promise((resolve) => setTimeout(resolve, 100)); + + yield { + type: 'user', + session_id: sessionId, + message: { + role: 'user', + content: 'What is 2 + 2?', + }, + parent_tool_use_id: null, + } as CLIUserMessage; + } + + const q = query({ + prompt: createMultiTurnConversation(), + options: { + ...SHARED_TEST_OPTIONS, + includePartialMessages: true, + debug: false, + }, + }); + + const messages: CLIMessage[] = []; + let partialMessageCount = 0; + let assistantMessageCount = 0; + + try { + for await (const message of q) { + messages.push(message); + + if (isCLIPartialAssistantMessage(message)) { + partialMessageCount++; + } + + if (isCLIAssistantMessage(message)) { + assistantMessageCount++; + } + } + + expect(messages.length).toBeGreaterThan(0); + expect(partialMessageCount).toBeGreaterThan(0); + expect(assistantMessageCount).toBeGreaterThanOrEqual(2); + } finally { + await q.close(); + } + }); + }); }); diff --git a/packages/sdk-typescript/test/e2e/single-turn.test.ts b/packages/sdk-typescript/test/e2e/single-turn.test.ts index 047be4f2..93c1ecc8 100644 --- a/packages/sdk-typescript/test/e2e/single-turn.test.ts +++ b/packages/sdk-typescript/test/e2e/single-turn.test.ts @@ -9,6 +9,7 @@ import { isCLIAssistantMessage, isCLISystemMessage, isCLIResultMessage, + isCLIPartialAssistantMessage, type TextBlock, type ContentBlock, type CLIMessage, @@ -327,6 +328,41 @@ describe('Single-Turn Query (E2E)', () => { await q.close(); } }); + + it('should receive partial messages when includePartialMessages is enabled', async () => { + const q = query({ + prompt: 'Count from 1 to 5', + options: { + ...SHARED_TEST_OPTIONS, + includePartialMessages: true, + debug: false, + }, + }); + + const messages: CLIMessage[] = []; + let partialMessageCount = 0; + let assistantMessageCount = 0; + + try { + for await (const message of q) { + messages.push(message); + + if (isCLIPartialAssistantMessage(message)) { + partialMessageCount++; + } + + if (isCLIAssistantMessage(message)) { + assistantMessageCount++; + } + } + + expect(messages.length).toBeGreaterThan(0); + expect(partialMessageCount).toBeGreaterThan(0); + expect(assistantMessageCount).toBeGreaterThan(0); + } finally { + await q.close(); + } + }); }); describe('Message Type Recognition', () => {