feat: add full message-level streaming + structured outputs (#37)

* feat: surface structured output in streams

Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>

* feat: simplify streaming result API

Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>

* test: move stress examples out of structured output PR

Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>

* fix: tighten streaming API types

Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>

---------

Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>

Author: varin-nair-factory

examples/readme-structured-output.ts

@@ -1,5 +1,7 @@
 import { OutputFormatType, run } from '@factory/droid-sdk';
 
+type FavoriteNumber = { favoriteNumber: number };
+
 const result = await run('Pick a favorite number between 1 and 42.', {
   cwd: process.cwd(),
   outputFormat: {
@@ -18,4 +20,6 @@ const result = await run('Pick a favorite number between 1 and 42.', {
   },
 });
 
-console.log(result.structuredOutput?.favoriteNumber);
+console.log(
+  (result.structuredOutput as FavoriteNumber | undefined)?.favoriteNumber
+);

examples/session-stream.ts

@@ -2,7 +2,7 @@
  * Simple session streaming example.
  *
  * Demonstrates using `session.stream()` to send a prompt, streaming
- * `AssistantTextDelta` messages to stdout, and handling `TurnComplete`.
+ * full assistant/tool messages, and handling the final `result`.
  *
  * Usage:
  *   npx tsx examples/session-stream.ts
@@ -20,19 +20,19 @@ async function main(): Promise<void> {
   try {
     for await (const msg of session.stream(prompt)) {
       switch (msg.type) {
-        case DroidMessageType.AssistantTextDelta:
+        case DroidMessageType.Assistant:
           process.stdout.write(msg.text);
           break;
 
-        case DroidMessageType.ToolUse:
-          console.log(`\n[Tool] ${msg.toolName}`);
+        case DroidMessageType.ToolCall:
+          console.log(`\n[Tool] ${msg.toolUse.name}`);
           break;
 
         case DroidMessageType.ToolResult:
           console.log(`[Tool Result] ${msg.isError ? 'Error' : 'OK'}`);
           break;
 
-        case DroidMessageType.TurnComplete:
+        case DroidMessageType.Result:
           console.log('\n\n--- Turn complete ---');
           if (msg.tokenUsage) {
             console.log(

examples/structured-output.ts

@@ -13,6 +13,8 @@ import assert from 'node:assert/strict';
 
 import { OutputFormatType, run } from '@factory/droid-sdk';
 
+type FavoriteNumber = { favoriteNumber: number };
+
 async function main(): Promise<void> {
   const prompt =
     process.argv.slice(2).join(' ') ||
@@ -40,11 +42,13 @@ async function main(): Promise<void> {
     outputFormat,
   });
 
-  assert.ok(result.structuredOutput, 'Expected structuredOutput to be set');
-  assert.equal(typeof result.structuredOutput['favoriteNumber'], 'number');
+  const structuredOutput = result.structuredOutput as FavoriteNumber | null;
+
+  assert.ok(structuredOutput, 'Expected structuredOutput to be set');
+  assert.equal(typeof structuredOutput.favoriteNumber, 'number');
 
   console.log('=== Structured output ===');
-  console.log(JSON.stringify(result.structuredOutput, null, 2));
+  console.log(JSON.stringify(structuredOutput, null, 2));
 
   console.log('\nStructured output example passed');
 }

examples/test-compact.ts

@@ -27,11 +27,13 @@ async function main(): Promise<void> {
       console.log(`=== Turn ${i + 1} ===`);
       console.log(`Prompt: "${prompts[i]}"\n`);
 
-      for await (const msg of session.stream(prompts[i])) {
+      for await (const msg of session.stream(prompts[i], {
+        includePartialMessages: true,
+      })) {
         if (msg.type === DroidMessageType.AssistantTextDelta) {
           process.stdout.write(msg.text);
         }
-        if (msg.type === DroidMessageType.TurnComplete) {
+        if (msg.type === DroidMessageType.Result) {
           console.log('\n');
         }
       }

src/helpers.ts

@@ -8,6 +8,7 @@ import type { DroidMcpServerConfig } from './mcp.js';
 import type {
   InitializeSessionRequestParams,
   McpServerConfig,
+  OutputFormat,
   SessionTag,
 } from './schemas/client.js';
 import type {
@@ -23,9 +24,11 @@ import type { ToolSelectionOverrides } from './schemas/shared.js';
 import {
   convertNotificationToStreamMessage,
   DroidMessageType,
+  isDefaultStreamMessage,
+  isInternalMessage,
   StreamStateTracker,
 } from './stream.js';
-import type { DroidMessage } from './stream.js';
+import type { DroidStreamEvent, InternalDroidMessage } from './stream.js';
 import { ProcessTransport } from './transport.js';
 import type { DroidClientTransport, ProcessTransportOptions } from './types.js';
 
@@ -52,14 +55,27 @@ export function extractInnerNotification(
 }
 
 export class MessageBridge {
-  private readonly _queue: DroidMessage[] = [];
+  private readonly _queue: DroidStreamEvent[] = [];
   private readonly _onDone: (() => void) | undefined;
   private _resolveWaiting: (() => void) | null = null;
   private _done = false;
-  private readonly _stateTracker = new StreamStateTracker();
+  private readonly _stateTracker: StreamStateTracker;
 
-  constructor(onDone?: () => void) {
+  constructor(
+    onDone?: () => void,
+    private readonly _options: {
+      includePartialMessages?: boolean;
+      sessionId?: string;
+      startedAt?: number;
+      outputFormat?: OutputFormat;
+    } = {}
+  ) {
     this._onDone = onDone;
+    this._stateTracker = new StreamStateTracker({
+      sessionId: _options.sessionId,
+      startedAt: _options.startedAt,
+      hasOutputFormat: _options.outputFormat !== undefined,
+    });
   }
 
   readonly notificationHandler = (
@@ -75,11 +91,15 @@ export class MessageBridge {
 
     for (const msg of messages) {
       const { message, additional } = this._stateTracker.processMessage(msg);
-      this._enqueue(message);
+      if (message && this._shouldYield(message)) {
+        this._enqueue(message);
+      }
 
       for (const extra of additional) {
-        this._enqueue(extra);
-        if (extra.type === DroidMessageType.TurnComplete) {
+        if (this._shouldYield(extra)) {
+          this._enqueue(extra);
+        }
+        if (extra.type === DroidMessageType.Result) {
           this._signalDone();
         }
       }
@@ -90,13 +110,13 @@ export class MessageBridge {
     this._signalDone();
   }
 
-  async *messages(): AsyncGenerator<DroidMessage, void, undefined> {
+  async *messages(): AsyncGenerator<DroidStreamEvent, void, undefined> {
     while (true) {
       while (this._queue.length > 0) {
         const msg = this._queue.shift()!;
         yield msg;
 
-        if (msg.type === DroidMessageType.TurnComplete) {
+        if (msg.type === DroidMessageType.Result) {
           return;
         }
       }
@@ -111,7 +131,18 @@ export class MessageBridge {
     }
   }
 
-  private _enqueue(msg: DroidMessage): void {
+  private _shouldYield(
+    message: InternalDroidMessage
+  ): message is DroidStreamEvent {
+    if (isInternalMessage(message)) {
+      return false;
+    }
+    return this._options.includePartialMessages
+      ? true
+      : isDefaultStreamMessage(message);
+  }
+
+  private _enqueue(msg: DroidStreamEvent): void {
     this._queue.push(msg);
     if (this._resolveWaiting) {
       const resolve = this._resolveWaiting;

src/index.ts

@@ -28,8 +28,19 @@ export {
   StreamStateTracker,
 } from './stream.js';
 export type {
+  DroidAssistantMessage,
+  DroidUserMessage,
+  DroidToolCallMessage,
+  DroidToolResultMessage,
+  DroidErrorMessage,
+  DroidResultMessage,
+  DroidStreamMessage,
+  DroidStreamEvent,
   AssistantTextDelta,
+  AssistantTextComplete,
   ThinkingTextDelta,
+  ThinkingTextComplete,
+  ToolCallDelta,
   ToolUse,
   ToolResult,
   ToolProgress,
@@ -48,8 +59,9 @@ export type {
   MissionWorkerCompleted,
   McpAuthRequired,
   McpAuthCompleted,
+  StructuredOutput,
+  StructuredOutputFields,
   ErrorEvent,
-  TurnComplete,
   DroidMessage,
 } from './stream.js';
 

src/run.ts

@@ -1,13 +1,15 @@
 import {
-  aggregateMessages,
   createSession,
   type CreateSessionOptions,
   type DroidResult,
   type MessageOptions,
 } from './session.js';
-import type { DroidMessage } from './stream.js';
+import { DroidMessageType } from './stream.js';
 
-export interface RunOptions extends CreateSessionOptions, MessageOptions {}
+export interface RunOptions
+  extends
+    CreateSessionOptions,
+    Omit<MessageOptions, 'includePartialMessages'> {}
 
 export async function run(
   prompt: string,
@@ -16,12 +18,16 @@ export async function run(
   const session = await createSession(options);
 
   try {
-    const startedAt = Date.now();
-    const messages: DroidMessage[] = [];
-    for await (const msg of session.stream(prompt, options)) {
-      messages.push(msg);
+    for await (const msg of session.stream(prompt, {
+      ...options,
+      includePartialMessages: false,
+    })) {
+      if (msg.type === DroidMessageType.Result) {
+        return msg;
+      }
     }
-    return aggregateMessages(session.sessionId, messages, startedAt, options);
+
+    throw new Error('Stream completed without a result message');
   } finally {
     await session.close();
   }

src/schemas/enums.ts

@@ -56,7 +56,10 @@ export enum SessionNotificationType {
   SESSION_TITLE_UPDATED = 'session_title_updated',
   MCP_STATUS_CHANGED = 'mcp_status_changed',
   ASSISTANT_TEXT_DELTA = 'assistant_text_delta',
+  ASSISTANT_TEXT_COMPLETE = 'assistant_text_complete',
   THINKING_TEXT_DELTA = 'thinking_text_delta',
+  THINKING_TEXT_COMPLETE = 'thinking_text_complete',
+  TOOL_CALL = 'tool_call',
   SESSION_TOKEN_USAGE_CHANGED = 'session_token_usage_changed',
   MISSION_STATE_CHANGED = 'mission_state_changed',
   MISSION_FEATURES_CHANGED = 'mission_features_changed',
@@ -66,6 +69,7 @@ export enum SessionNotificationType {
   MISSION_WORKER_COMPLETED = 'mission_worker_completed',
   MCP_AUTH_REQUIRED = 'mcp_auth_required',
   MCP_AUTH_COMPLETED = 'mcp_auth_completed',
+  STRUCTURED_OUTPUT = 'structured_output',
 }
 
 /** Tool confirmation outcome options (possible user responses to permission requests). */

src/schemas/index.ts

@@ -510,7 +510,12 @@ export {
   SettingsUpdatedNotificationSchema,
   SettingsUpdatedPayloadSchema,
   StartMissionRunConfirmationDetailsSchema,
+  AssistantTextCompleteNotificationSchema,
+  StructuredOutputErrorSchema,
+  StructuredOutputNotificationSchema,
+  ThinkingTextCompleteNotificationSchema,
   ThinkingTextDeltaNotificationSchema,
+  ToolCallNotificationSchema,
   ToolConfirmationDetailsSchema,
   ToolConfirmationInfoSchema,
   ToolProgressUpdateNotificationSchema,
@@ -528,6 +533,7 @@ export type {
   AskUserResponse,
   AskUserResult,
   ApplyPatchToolConfirmationDetails,
+  AssistantTextCompleteNotification,
   AssistantTextDeltaNotification,
   CliRequestOrNotification,
   CreateMessageNotification,
@@ -564,7 +570,11 @@ export type {
   SettingsUpdatedNotification,
   SettingsUpdatedPayload,
   StartMissionRunConfirmationDetails,
+  StructuredOutputError,
+  StructuredOutputNotification,
+  ThinkingTextCompleteNotification,
   ThinkingTextDeltaNotification,
+  ToolCallNotification,
   ToolConfirmationDetails,
   ToolConfirmationInfo,
   ToolProgressUpdate,