diff --git a/src/helpers.ts b/src/helpers.ts index 9642de3..8b469f5 100644 --- a/src/helpers.ts +++ b/src/helpers.ts @@ -84,7 +84,10 @@ export class MessageBridge { const inner = extractInnerNotification(notification); if (!inner) return; - const converted = convertNotificationToStreamMessage(inner); + const converted = convertNotificationToStreamMessage( + inner, + this._options.sessionId ?? '' + ); if (converted === null) return; const messages = Array.isArray(converted) ? converted : [converted]; diff --git a/src/stream.ts b/src/stream.ts index d1af7d7..85e7688 100644 --- a/src/stream.ts +++ b/src/stream.ts @@ -59,62 +59,66 @@ export const DroidMessageType = { Result: 'result', } as const; -export interface AssistantTextDelta { +export interface DroidStreamBase { + readonly sessionId: string; +} + +export interface AssistantTextDelta extends DroidStreamBase { readonly type: 'assistant_text_delta'; readonly messageId: string; readonly blockIndex: number; readonly text: string; } -export interface AssistantTextComplete { +export interface AssistantTextComplete extends DroidStreamBase { readonly type: 'assistant_text_complete'; readonly messageId: string; readonly blockIndex: number; } -export interface ThinkingTextDelta { +export interface ThinkingTextDelta extends DroidStreamBase { readonly type: 'thinking_text_delta'; readonly messageId: string; readonly blockIndex: number; readonly text: string; } -export interface ThinkingTextComplete { +export interface ThinkingTextComplete extends DroidStreamBase { readonly type: 'thinking_text_complete'; readonly messageId: string; readonly blockIndex: number; readonly durationMs?: number; } -export interface ToolCallDelta { +export interface ToolCallDelta extends DroidStreamBase { readonly type: 'tool_call_delta'; readonly toolUse: ToolUseBlock; } -export interface ToolUse { +export interface ToolUse extends DroidStreamBase { readonly type: 'tool_use'; readonly toolName: string; readonly toolInput: JsonObject; readonly toolUseId: string; } -export interface DroidToolCallMessage { +export interface DroidToolCallMessage extends DroidStreamBase { readonly type: 'tool_call'; readonly toolUse: ToolUseBlock; } -export interface DroidAssistantMessage { +export interface DroidAssistantMessage extends DroidStreamBase { readonly type: 'assistant'; readonly message: FactoryDroidMessage; readonly text: string; } -export interface DroidUserMessage { +export interface DroidUserMessage extends DroidStreamBase { readonly type: 'user'; readonly message: FactoryDroidMessage; } -export interface ToolResult { +export interface ToolResult extends DroidStreamBase { readonly type: 'tool_result'; readonly toolUseId: string; readonly toolName: string; @@ -122,7 +126,7 @@ export interface ToolResult { readonly isError: boolean; } -export interface ToolProgress { +export interface ToolProgress extends DroidStreamBase { readonly type: 'tool_progress'; readonly toolUseId: string; readonly toolName: string; @@ -130,7 +134,7 @@ export interface ToolProgress { readonly update: ToolProgressUpdate; } -export interface WorkingStateChanged { +export interface WorkingStateChanged extends DroidStreamBase { readonly type: 'working_state_changed'; readonly state: DroidWorkingState; } @@ -138,10 +142,11 @@ export interface WorkingStateChanged { export type TokenUsageUpdate = Readonly< { type: 'token_usage_update'; - } & TokenUsage + } & DroidStreamBase & + TokenUsage >; -export interface CreateMessage { +export interface CreateMessage extends DroidStreamBase { readonly type: 'create_message'; readonly messageId: CreateMessageNotification['message']['id']; readonly role: CreateMessageNotification['message']['role']; @@ -149,61 +154,61 @@ export interface CreateMessage { readonly parentId?: CreateMessageNotification['parentId']; } -export interface PermissionResolved { +export interface PermissionResolved extends DroidStreamBase { readonly type: 'permission_resolved'; readonly requestId: string; readonly toolUseIds: string[]; readonly selectedOption: ToolConfirmationOutcome; } -export interface SettingsUpdated { +export interface SettingsUpdated extends DroidStreamBase { readonly type: 'settings_updated'; readonly settings: SettingsUpdatedPayload; } -export interface SessionTitleUpdated { +export interface SessionTitleUpdated extends DroidStreamBase { readonly type: 'session_title_updated'; readonly title: string; } -export interface McpStatusChanged { +export interface McpStatusChanged extends DroidStreamBase { readonly type: 'mcp_status_changed'; readonly servers: McpServerStatusInfo[]; readonly summary: McpStatusSummary; } -export interface MissionStateChanged { +export interface MissionStateChanged extends DroidStreamBase { readonly type: 'mission_state_changed'; readonly state: MissionState; } -export interface MissionFeaturesChanged { +export interface MissionFeaturesChanged extends DroidStreamBase { readonly type: 'mission_features_changed'; readonly features: MissionFeature[]; } -export interface MissionProgressEntry { +export interface MissionProgressEntry extends DroidStreamBase { readonly type: 'mission_progress_entry'; readonly progressLog: ProgressLogEntry[]; } -export interface MissionHeartbeat { +export interface MissionHeartbeat extends DroidStreamBase { readonly type: 'mission_heartbeat'; readonly timestamp: string; } -export interface MissionWorkerStarted { +export interface MissionWorkerStarted extends DroidStreamBase { readonly type: 'mission_worker_started'; readonly workerSessionId: string; } -export interface MissionWorkerCompleted { +export interface MissionWorkerCompleted extends DroidStreamBase { readonly type: 'mission_worker_completed'; readonly workerSessionId: string; readonly exitCode: number; } -export interface McpAuthRequired { +export interface McpAuthRequired extends DroidStreamBase { readonly type: 'mcp_auth_required'; readonly serverName: string; readonly authUrl: string; @@ -211,14 +216,14 @@ export interface McpAuthRequired { readonly state: string; } -export interface McpAuthCompleted { +export interface McpAuthCompleted extends DroidStreamBase { readonly type: 'mcp_auth_completed'; readonly serverName: string; readonly outcome: McpAuthOutcome; readonly message: string; } -export interface HookExecution { +export interface HookExecution extends DroidStreamBase { readonly type: 'hook'; readonly hookId: string; readonly eventName?: DroidHookEvent; @@ -237,12 +242,13 @@ export interface StructuredOutputFields { readonly structuredOutputError: ServerStructuredOutputError | null; } -export interface StructuredOutput extends StructuredOutputFields { +export interface StructuredOutput + extends DroidStreamBase, StructuredOutputFields { readonly type: 'structured_output'; readonly messageId: string; } -export interface ErrorEvent { +export interface ErrorEvent extends DroidStreamBase { readonly type: 'error'; readonly message: string; readonly errorType: ErrorNotification['errorType']; @@ -258,9 +264,8 @@ export type DroidResultSubtype = | 'error_during_execution' | 'error_structured_output'; -interface DroidResultBase { +interface DroidResultBase extends DroidStreamBase { readonly type: 'result'; - readonly sessionId: string; readonly durationMs: number; readonly numTurns: number; readonly result: string; @@ -338,7 +343,8 @@ export type DroidMessageType = (typeof DroidMessageType)[keyof typeof DroidMessageType]; export function convertNotificationToStreamMessage( - raw: unknown + raw: unknown, + sessionId: string ): InternalDroidMessage | InternalDroidMessage[] | null { const parsed = SessionNotificationPayloadSchema.safeParse(raw); if (!parsed.success) { @@ -351,6 +357,7 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.ASSISTANT_TEXT_DELTA: return { type: DroidMessageType.AssistantTextDelta, + sessionId, messageId: notification.messageId, blockIndex: notification.blockIndex, text: notification.textDelta, @@ -359,6 +366,7 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.ASSISTANT_TEXT_COMPLETE: return { type: DroidMessageType.AssistantTextComplete, + sessionId, messageId: notification.messageId, blockIndex: notification.blockIndex, }; @@ -366,6 +374,7 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.THINKING_TEXT_DELTA: return { type: DroidMessageType.ThinkingTextDelta, + sessionId, messageId: notification.messageId, blockIndex: notification.blockIndex, text: notification.textDelta, @@ -374,6 +383,7 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.THINKING_TEXT_COMPLETE: return { type: DroidMessageType.ThinkingTextComplete, + sessionId, messageId: notification.messageId, blockIndex: notification.blockIndex, durationMs: notification.durationMs, @@ -382,12 +392,14 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.TOOL_CALL: return { type: DroidMessageType.ToolCallDelta, + sessionId, toolUse: notification.toolUse, }; case SessionNotificationType.TOOL_RESULT: return { type: DroidMessageType.ToolResult, + sessionId, toolUseId: notification.toolUseId, toolName: '', content: normalizeToolResultContent(notification.content), @@ -399,6 +411,7 @@ export function convertNotificationToStreamMessage( const text = update?.text ?? update?.status ?? update?.details ?? ''; return { type: DroidMessageType.ToolProgress, + sessionId, toolUseId: notification.toolUseId, toolName: notification.toolName, content: text, @@ -409,6 +422,7 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.DROID_WORKING_STATE_CHANGED: return { type: DroidMessageType.WorkingStateChanged, + sessionId, state: notification.newState, }; @@ -416,6 +430,7 @@ export function convertNotificationToStreamMessage( const tu: TokenUsage = notification.tokenUsage; return { type: DroidMessageType.TokenUsageUpdate, + sessionId, inputTokens: tu.inputTokens, outputTokens: tu.outputTokens, cacheReadTokens: tu.cacheReadTokens, @@ -432,6 +447,7 @@ export function convertNotificationToStreamMessage( if (block.type === 'tool_use') { messages.push({ type: DroidMessageType.ToolCall, + sessionId, toolUse: block, }); } @@ -440,17 +456,20 @@ export function convertNotificationToStreamMessage( if (msg.role === FactoryDroidMessageRole.Assistant) { messages.push({ type: DroidMessageType.Assistant, + sessionId, message: msg, text: extractTextFromMessage(msg), }); } else if (msg.role === FactoryDroidMessageRole.User) { messages.push({ type: DroidMessageType.User, + sessionId, message: msg, }); } else { messages.push({ type: 'create_message', + sessionId, messageId: msg.id, role: msg.role, content: msg.content, @@ -464,6 +483,7 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.ERROR: return { type: DroidMessageType.Error, + sessionId, message: notification.message, errorType: notification.errorType, timestamp: notification.timestamp, @@ -472,6 +492,7 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.PERMISSION_RESOLVED: return { type: DroidMessageType.PermissionResolved, + sessionId, requestId: notification.requestId, toolUseIds: notification.toolUseIds, selectedOption: notification.selectedOption, @@ -480,18 +501,21 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.SETTINGS_UPDATED: return { type: DroidMessageType.SettingsUpdated, + sessionId, settings: notification.settings, }; case SessionNotificationType.SESSION_TITLE_UPDATED: return { type: DroidMessageType.SessionTitleUpdated, + sessionId, title: notification.title, }; case SessionNotificationType.MCP_STATUS_CHANGED: return { type: DroidMessageType.McpStatusChanged, + sessionId, servers: notification.servers, summary: notification.summary, }; @@ -499,36 +523,42 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.MISSION_STATE_CHANGED: return { type: DroidMessageType.MissionStateChanged, + sessionId, state: notification.state, }; case SessionNotificationType.MISSION_FEATURES_CHANGED: return { type: DroidMessageType.MissionFeaturesChanged, + sessionId, features: notification.features, }; case SessionNotificationType.MISSION_PROGRESS_ENTRY: return { type: DroidMessageType.MissionProgressEntry, + sessionId, progressLog: notification.progressLog, }; case SessionNotificationType.MISSION_HEARTBEAT: return { type: DroidMessageType.MissionHeartbeat, + sessionId, timestamp: notification.timestamp, }; case SessionNotificationType.MISSION_WORKER_STARTED: return { type: DroidMessageType.MissionWorkerStarted, + sessionId, workerSessionId: notification.workerSessionId, }; case SessionNotificationType.MISSION_WORKER_COMPLETED: return { type: DroidMessageType.MissionWorkerCompleted, + sessionId, workerSessionId: notification.workerSessionId, exitCode: notification.exitCode, }; @@ -536,6 +566,7 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.MCP_AUTH_REQUIRED: return { type: DroidMessageType.McpAuthRequired, + sessionId, serverName: notification.serverName, authUrl: notification.authUrl, message: notification.message, @@ -545,6 +576,7 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.MCP_AUTH_COMPLETED: return { type: DroidMessageType.McpAuthCompleted, + sessionId, serverName: notification.serverName, outcome: notification.outcome, message: notification.message, @@ -553,6 +585,7 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.HOOK_EXECUTION_STARTED: return notification.hookCommands.map((hookCommand) => ({ type: DroidMessageType.Hook, + sessionId, hookId: notification.hookId, eventName: notification.hookEventName, matcher: notification.hookMatcher, @@ -565,6 +598,7 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.HOOK_EXECUTION_COMPLETED: return (notification.hookResults ?? []).map((hookResult) => ({ type: DroidMessageType.Hook, + sessionId, hookId: notification.hookId, eventName: notification.hookEventName, matcher: notification.hookMatcher, @@ -580,6 +614,7 @@ export function convertNotificationToStreamMessage( case SessionNotificationType.STRUCTURED_OUTPUT: return { type: 'structured_output', + sessionId, messageId: notification.messageId, structuredOutput: notification.structuredOutput, structuredOutputError: notification.structuredOutputError, diff --git a/tests/helpers.test.ts b/tests/helpers.test.ts index bbd818c..2792f34 100644 --- a/tests/helpers.test.ts +++ b/tests/helpers.test.ts @@ -168,6 +168,31 @@ describe('MessageBridge', () => { expect(messages.some((m) => m.type === 'working_state_changed')).toBe(true); }); + it('passes configured sessionId to converted messages', async () => { + bridge = new MessageBridge(undefined, { + includePartialMessages: true, + sessionId: 'sess-bridge', + }); + bridge.notificationHandler( + makeSessionNotification( + SessionNotificationType.DROID_WORKING_STATE_CHANGED, + { newState: DroidWorkingState.StreamingAssistantMessage } + ) as Record + ); + + bridge.signalDone(); + + const messages = []; + for await (const msg of bridge.messages()) { + messages.push(msg); + } + + expect(messages[0]).toMatchObject({ + type: 'working_state_changed', + sessionId: 'sess-bridge', + }); + }); + it('terminates generator on result message', async () => { // Transition to streaming state then back to idle to trigger result bridge.notificationHandler( diff --git a/tests/integration.test.ts b/tests/integration.test.ts index 767dd54..8f252db 100644 --- a/tests/integration.test.ts +++ b/tests/integration.test.ts @@ -1808,7 +1808,8 @@ describe('Settings update notification flow (VAL-CROSS-007)', () => { | undefined; if (inner) { const converted = convertNotificationToStreamMessage( - inner as { type: string; [key: string]: unknown } + inner as { type: string; [key: string]: unknown }, + session.sessionId ); if ( converted && diff --git a/tests/run.test.ts b/tests/run.test.ts index 06e1198..ac8bcf0 100644 --- a/tests/run.test.ts +++ b/tests/run.test.ts @@ -290,6 +290,7 @@ describe('run()', () => { expect(result.structuredOutput).toEqual({ name: 'Ada' }); expect(result.messages).toContainEqual({ type: 'assistant', + sessionId: 'sess-run-structured-output', text: JSON.stringify({ name: 'Ada' }), message: expect.objectContaining({ id: 'msg-structured', diff --git a/tests/stream.test.ts b/tests/stream.test.ts index bd4c5ee..6070cf0 100644 --- a/tests/stream.test.ts +++ b/tests/stream.test.ts @@ -12,7 +12,7 @@ import { ToolConfirmationOutcome, } from '../src/schemas/index.js'; import { - convertNotificationToStreamMessage, + convertNotificationToStreamMessage as convertNotificationToStreamMessageRaw, DroidMessageType, StreamStateTracker, } from '../src/stream.js'; @@ -48,6 +48,13 @@ function makeNotification(type: string, payload: Record) { return { type, ...payload }; } +function convertNotificationToStreamMessage( + raw: unknown, + sessionId = 'sess-test' +) { + return convertNotificationToStreamMessageRaw(raw, sessionId); +} + const expectedDroidMessageTypes = [ 'assistant', 'user', @@ -96,6 +103,7 @@ describe('DroidMessage types', () => { it('AssistantTextDelta has correct structure', () => { const msg: AssistantTextDelta = { type: 'assistant_text_delta', + sessionId: 's1', messageId: 'msg-1', blockIndex: 0, text: 'Hello', @@ -109,6 +117,7 @@ describe('DroidMessage types', () => { it('ThinkingTextDelta has correct structure', () => { const msg: ThinkingTextDelta = { type: 'thinking_text_delta', + sessionId: 's1', messageId: 'msg-2', blockIndex: 1, text: 'Let me think...', @@ -120,6 +129,7 @@ describe('DroidMessage types', () => { it('ToolUse has correct structure', () => { const msg: ToolUse = { type: 'tool_use', + sessionId: 's1', toolName: 'read_file', toolInput: { path: '/tmp/test.txt' }, toolUseId: 'tu-1', @@ -133,6 +143,7 @@ describe('DroidMessage types', () => { it('ToolResult has correct structure', () => { const msg: ToolResult = { type: 'tool_result', + sessionId: 's1', toolUseId: 'tu-1', toolName: 'read_file', content: 'file contents here', @@ -147,6 +158,7 @@ describe('DroidMessage types', () => { it('ToolProgress has correct structure', () => { const msg: ToolProgress = { type: 'tool_progress', + sessionId: 's1', toolUseId: 'tu-1', toolName: 'execute', content: 'Running...', @@ -160,6 +172,7 @@ describe('DroidMessage types', () => { it('WorkingStateChanged has correct structure', () => { const msg: WorkingStateChanged = { type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.ExecutingTool, }; expect(msg.type).toBe('working_state_changed'); @@ -169,6 +182,7 @@ describe('DroidMessage types', () => { it('TokenUsageUpdate has correct structure', () => { const msg: TokenUsageUpdate = { type: 'token_usage_update', + sessionId: 's1', inputTokens: 100, outputTokens: 50, cacheReadTokens: 10, @@ -186,6 +200,7 @@ describe('DroidMessage types', () => { it('CreateMessage has correct structure', () => { const msg: CreateMessage = { type: 'create_message', + sessionId: 's1', messageId: 'msg-3', role: 'assistant', content: [{ type: 'text', text: 'hi' }], @@ -200,6 +215,7 @@ describe('DroidMessage types', () => { it('PermissionResolved has correct structure', () => { const msg: PermissionResolved = { type: 'permission_resolved', + sessionId: 's1', requestId: 'req-1', toolUseIds: ['tu-1', 'tu-2'], selectedOption: ToolConfirmationOutcome.ProceedOnce, @@ -211,6 +227,7 @@ describe('DroidMessage types', () => { it('SettingsUpdated has correct structure', () => { const msg: SettingsUpdated = { type: 'settings_updated', + sessionId: 's1', settings: { modelId: 'claude-opus-4' }, }; expect(msg.type).toBe('settings_updated'); @@ -220,6 +237,7 @@ describe('DroidMessage types', () => { it('SessionTitleUpdated has correct structure', () => { const msg: SessionTitleUpdated = { type: 'session_title_updated', + sessionId: 's1', title: 'My Session', }; expect(msg.type).toBe('session_title_updated'); @@ -229,6 +247,7 @@ describe('DroidMessage types', () => { it('McpStatusChanged has correct structure', () => { const msg: McpStatusChanged = { type: 'mcp_status_changed', + sessionId: 's1', servers: [ { name: 'test-server', @@ -246,6 +265,7 @@ describe('DroidMessage types', () => { it('MissionStateChanged has correct structure', () => { const msg: MissionStateChanged = { type: 'mission_state_changed', + sessionId: 's1', state: MissionState.Running, }; expect(msg.type).toBe('mission_state_changed'); @@ -255,6 +275,7 @@ describe('DroidMessage types', () => { it('MissionFeaturesChanged has correct structure', () => { const msg: MissionFeaturesChanged = { type: 'mission_features_changed', + sessionId: 's1', features: [ { id: 'feat-1', @@ -274,6 +295,7 @@ describe('DroidMessage types', () => { it('MissionProgressEntry has correct structure', () => { const msg: MissionProgressEntry = { type: 'mission_progress_entry', + sessionId: 's1', progressLog: [], }; expect(msg.type).toBe('mission_progress_entry'); @@ -283,6 +305,7 @@ describe('DroidMessage types', () => { it('MissionHeartbeat has correct structure', () => { const msg: MissionHeartbeat = { type: 'mission_heartbeat', + sessionId: 's1', timestamp: '2025-01-01T00:00:00Z', }; expect(msg.type).toBe('mission_heartbeat'); @@ -292,6 +315,7 @@ describe('DroidMessage types', () => { it('MissionWorkerStarted has correct structure', () => { const msg: MissionWorkerStarted = { type: 'mission_worker_started', + sessionId: 's1', workerSessionId: 'ws-1', }; expect(msg.type).toBe('mission_worker_started'); @@ -301,6 +325,7 @@ describe('DroidMessage types', () => { it('MissionWorkerCompleted has correct structure', () => { const msg: MissionWorkerCompleted = { type: 'mission_worker_completed', + sessionId: 's1', workerSessionId: 'ws-1', exitCode: 0, }; @@ -311,6 +336,7 @@ describe('DroidMessage types', () => { it('McpAuthRequired has correct structure', () => { const msg: McpAuthRequired = { type: 'mcp_auth_required', + sessionId: 's1', serverName: 'my-server', authUrl: 'https://auth.example.com', message: 'Please authenticate', @@ -323,6 +349,7 @@ describe('DroidMessage types', () => { it('McpAuthCompleted has correct structure', () => { const msg: McpAuthCompleted = { type: 'mcp_auth_completed', + sessionId: 's1', serverName: 'my-server', outcome: McpAuthOutcome.Success, message: 'Authenticated', @@ -334,6 +361,7 @@ describe('DroidMessage types', () => { it('HookExecution has correct structure', () => { const msg: HookExecution = { type: 'hook', + sessionId: 's1', hookId: 'hook-1', eventName: 'PreToolUse', matcher: 'Execute', @@ -348,6 +376,7 @@ describe('DroidMessage types', () => { it('ErrorEvent has correct structure', () => { const msg: ErrorEvent = { type: 'error', + sessionId: 's1', message: 'Something went wrong', errorType: DroidErrorType.SESSION_ERROR, timestamp: '2025-01-01T00:00:00Z', @@ -360,6 +389,7 @@ describe('DroidMessage types', () => { it('DroidResultMessage has correct structure', () => { const tokenUsage: TokenUsageUpdate = { type: 'token_usage_update', + sessionId: 's1', inputTokens: 100, outputTokens: 50, cacheReadTokens: 10, @@ -390,6 +420,7 @@ describe('DroidMessage types', () => { const messages: DroidMessage[] = [ { type: 'assistant', + sessionId: 's1', message: { id: 'a1', role: 'assistant', @@ -401,6 +432,7 @@ describe('DroidMessage types', () => { }, { type: 'user', + sessionId: 's1', message: { id: 'u1', role: 'user', @@ -411,6 +443,7 @@ describe('DroidMessage types', () => { }, { type: 'tool_call', + sessionId: 's1', toolUse: { type: 'tool_use', id: 'tu-1', @@ -420,28 +453,33 @@ describe('DroidMessage types', () => { }, { type: 'assistant_text_delta', + sessionId: 's1', messageId: 'm1', blockIndex: 0, text: 'hi', }, { type: 'assistant_text_complete', + sessionId: 's1', messageId: 'm1', blockIndex: 0, }, { type: 'thinking_text_delta', + sessionId: 's1', messageId: 'm1', blockIndex: 0, text: 'hmm', }, { type: 'thinking_text_complete', + sessionId: 's1', messageId: 'm1', blockIndex: 0, }, { type: 'tool_call_delta', + sessionId: 's1', toolUse: { type: 'tool_use', id: 'tu-1', @@ -451,6 +489,7 @@ describe('DroidMessage types', () => { }, { type: 'tool_result', + sessionId: 's1', toolUseId: 'tu1', toolName: 'x', content: '', @@ -458,14 +497,20 @@ describe('DroidMessage types', () => { }, { type: 'tool_progress', + sessionId: 's1', toolUseId: 'tu1', toolName: 'x', content: '...', update: { type: 'status' }, }, - { type: 'working_state_changed', state: DroidWorkingState.Idle }, + { + type: 'working_state_changed', + sessionId: 's1', + state: DroidWorkingState.Idle, + }, { type: 'token_usage_update', + sessionId: 's1', inputTokens: 0, outputTokens: 0, cacheReadTokens: 0, @@ -474,25 +519,41 @@ describe('DroidMessage types', () => { }, { type: 'permission_resolved', + sessionId: 's1', requestId: 'r1', toolUseIds: [], selectedOption: ToolConfirmationOutcome.Cancel, }, - { type: 'settings_updated', settings: {} }, - { type: 'session_title_updated', title: 't' }, + { type: 'settings_updated', sessionId: 's1', settings: {} }, + { type: 'session_title_updated', sessionId: 's1', title: 't' }, { type: 'mcp_status_changed', + sessionId: 's1', servers: [], summary: { total: 0, connected: 0, connecting: 0, failed: 0 }, }, - { type: 'mission_state_changed', state: MissionState.Running }, - { type: 'mission_features_changed', features: [] }, - { type: 'mission_progress_entry', progressLog: [] }, - { type: 'mission_heartbeat', timestamp: 't' }, - { type: 'mission_worker_started', workerSessionId: 'ws1' }, - { type: 'mission_worker_completed', workerSessionId: 'ws1', exitCode: 0 }, + { + type: 'mission_state_changed', + sessionId: 's1', + state: MissionState.Running, + }, + { type: 'mission_features_changed', sessionId: 's1', features: [] }, + { type: 'mission_progress_entry', sessionId: 's1', progressLog: [] }, + { type: 'mission_heartbeat', sessionId: 's1', timestamp: 't' }, + { + type: 'mission_worker_started', + sessionId: 's1', + workerSessionId: 'ws1', + }, + { + type: 'mission_worker_completed', + sessionId: 's1', + workerSessionId: 'ws1', + exitCode: 0, + }, { type: 'mcp_auth_required', + sessionId: 's1', serverName: 's', authUrl: 'u', message: 'm', @@ -500,12 +561,14 @@ describe('DroidMessage types', () => { }, { type: 'mcp_auth_completed', + sessionId: 's1', serverName: 's', outcome: McpAuthOutcome.Success, message: 'm', }, { type: 'hook', + sessionId: 's1', hookId: 'hook-1', eventName: 'PreToolUse', matcher: 'Execute', @@ -519,6 +582,7 @@ describe('DroidMessage types', () => { }, { type: 'error', + sessionId: 's1', message: 'err', errorType: DroidErrorType.ERROR, timestamp: 't', @@ -1161,6 +1225,7 @@ describe('convertNotificationToStreamMessage', () => { expect(result).toEqual([ { type: 'hook', + sessionId: 'sess-test', hookId: 'hook-1', eventName: 'PreToolUse', matcher: 'Execute', @@ -1171,6 +1236,7 @@ describe('convertNotificationToStreamMessage', () => { }, { type: 'hook', + sessionId: 'sess-test', hookId: 'hook-1', eventName: 'PreToolUse', matcher: 'Execute', @@ -1209,6 +1275,7 @@ describe('convertNotificationToStreamMessage', () => { expect(result).toEqual([ { type: 'hook', + sessionId: 'sess-test', hookId: 'hook-1', eventName: 'PreToolUse', matcher: 'Execute', @@ -1240,6 +1307,7 @@ describe('convertNotificationToStreamMessage', () => { expect(result).toEqual({ type: 'structured_output', + sessionId: 'sess-test', messageId: 'msg-structured', structuredOutput: { name: 'Ada' }, structuredOutputError: null, @@ -1264,6 +1332,7 @@ describe('convertNotificationToStreamMessage', () => { expect(result).toEqual({ type: 'structured_output', + sessionId: 'sess-test', messageId: 'msg-structured', structuredOutput: null, structuredOutputError: { @@ -1299,8 +1368,9 @@ describe('convertNotificationToStreamMessage', () => { ); }); - it('every notification type returns a non-null result (with valid payloads)', () => { + it('every notification type returns a non-null result with sessionId', () => { const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}); + const sessionId = 'sess-all-notifications'; const payloads: Record> = { [SessionNotificationType.ASSISTANT_TEXT_DELTA]: { @@ -1427,11 +1497,23 @@ describe('convertNotificationToStreamMessage', () => { const payload = payloads[notifType]; expect(payload, `Missing test payload for ${notifType}`).toBeDefined(); const notification = makeNotification(notifType, payload); - const result = convertNotificationToStreamMessage(notification); + const result = convertNotificationToStreamMessage( + notification, + sessionId + ); expect( result, `Converter returned null for ${notifType}` ).not.toBeNull(); + const messages = Array.isArray(result) ? result : [result!]; + expect(messages.length, `No messages for ${notifType}`).toBeGreaterThan( + 0 + ); + for (const message of messages) { + expect(message.sessionId, `Missing sessionId for ${notifType}`).toBe( + sessionId + ); + } } expect(warnSpy).not.toHaveBeenCalled(); @@ -1451,12 +1533,14 @@ describe('StreamStateTracker', () => { it('emits Result on non-idle → idle transition', () => { const r1 = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); expect(r1.additional).toEqual([]); const r2 = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(r2.additional).toHaveLength(1); @@ -1467,6 +1551,7 @@ describe('StreamStateTracker', () => { it('does NOT emit Result for initial idle', () => { const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(result.additional).toEqual([]); @@ -1475,10 +1560,12 @@ describe('StreamStateTracker', () => { it('does NOT emit Result for non-idle → non-idle transitions', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.ExecutingTool, }); expect(result.additional).toEqual([]); @@ -1487,14 +1574,17 @@ describe('StreamStateTracker', () => { it('emits Result after multiple non-idle states', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.ExecutingTool, }); const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(result.additional).toHaveLength(1); @@ -1504,10 +1594,12 @@ describe('StreamStateTracker', () => { it('can emit Result again after reset', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); @@ -1515,16 +1607,19 @@ describe('StreamStateTracker', () => { const r1 = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(r1.additional).toEqual([]); tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.ExecutingTool, }); const r2 = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(r2.additional).toHaveLength(1); @@ -1536,10 +1631,12 @@ describe('StreamStateTracker', () => { it('attaches structured output to Result', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); tracker.processMessage({ type: 'structured_output', + sessionId: 's1', messageId: 'msg-structured', structuredOutput: { name: 'Ada' }, structuredOutputError: null, @@ -1547,6 +1644,7 @@ describe('StreamStateTracker', () => { const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); @@ -1562,10 +1660,12 @@ describe('StreamStateTracker', () => { it('attaches structured output errors to Result', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); tracker.processMessage({ type: 'structured_output', + sessionId: 's1', messageId: 'msg-structured', structuredOutput: null, structuredOutputError: { @@ -1576,6 +1676,7 @@ describe('StreamStateTracker', () => { const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); @@ -1594,25 +1695,30 @@ describe('StreamStateTracker', () => { it('does not carry structured output across turns', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); tracker.processMessage({ type: 'structured_output', + sessionId: 's1', messageId: 'msg-structured', structuredOutput: { name: 'Ada' }, structuredOutputError: null, }); tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); @@ -1630,6 +1736,7 @@ describe('StreamStateTracker', () => { it('carries last-seen TokenUsageUpdate in Result', () => { const tokenUsage: TokenUsageUpdate = { type: 'token_usage_update', + sessionId: 's1', inputTokens: 100, outputTokens: 50, cacheReadTokens: 10, @@ -1639,11 +1746,13 @@ describe('StreamStateTracker', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); tracker.processMessage(tokenUsage); const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); @@ -1661,10 +1770,12 @@ describe('StreamStateTracker', () => { it('uses the LAST token usage update when multiple are received', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); tracker.processMessage({ type: 'token_usage_update', + sessionId: 's1', inputTokens: 50, outputTokens: 25, cacheReadTokens: 5, @@ -1673,6 +1784,7 @@ describe('StreamStateTracker', () => { }); tracker.processMessage({ type: 'token_usage_update', + sessionId: 's1', inputTokens: 200, outputTokens: 100, cacheReadTokens: 20, @@ -1682,6 +1794,7 @@ describe('StreamStateTracker', () => { const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); @@ -1693,10 +1806,12 @@ describe('StreamStateTracker', () => { it('returns null tokenUsage when no TokenUsageUpdate was received', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); @@ -1707,10 +1822,12 @@ describe('StreamStateTracker', () => { it('token usage is reset after reset()', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); tracker.processMessage({ type: 'token_usage_update', + sessionId: 's1', inputTokens: 100, outputTokens: 50, cacheReadTokens: 10, @@ -1719,6 +1836,7 @@ describe('StreamStateTracker', () => { }); tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); @@ -1726,10 +1844,12 @@ describe('StreamStateTracker', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.ExecutingTool, }); const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); @@ -1742,12 +1862,14 @@ describe('StreamStateTracker', () => { it('enriches tool_result with toolName from prior tool_use', () => { tracker.processMessage({ type: 'tool_use', + sessionId: 's1', toolName: 'read_file', toolInput: { path: '/tmp/test' }, toolUseId: 'tu-1', }); const { message } = tracker.processMessage({ type: 'tool_result', + sessionId: 's1', toolUseId: 'tu-1', toolName: '', content: 'file contents', @@ -1760,18 +1882,21 @@ describe('StreamStateTracker', () => { it('enriches tool_results from multiple tool_use mappings', () => { tracker.processMessage({ type: 'tool_use', + sessionId: 's1', toolName: 'read_file', toolInput: {}, toolUseId: 'tu-1', }); tracker.processMessage({ type: 'tool_use', + sessionId: 's1', toolName: 'write_file', toolInput: {}, toolUseId: 'tu-2', }); const r1 = tracker.processMessage({ type: 'tool_result', + sessionId: 's1', toolUseId: 'tu-1', toolName: '', content: '', @@ -1779,6 +1904,7 @@ describe('StreamStateTracker', () => { }); const r2 = tracker.processMessage({ type: 'tool_result', + sessionId: 's1', toolUseId: 'tu-2', toolName: '', content: '', @@ -1791,6 +1917,7 @@ describe('StreamStateTracker', () => { it('returns empty string toolName for unknown toolUseId', () => { const { message } = tracker.processMessage({ type: 'tool_result', + sessionId: 's1', toolUseId: 'unknown-id', toolName: '', content: '', @@ -1804,6 +1931,7 @@ describe('StreamStateTracker', () => { it('returns empty additional array for non-state-change messages', () => { const textDelta: DroidMessage = { type: 'assistant_text_delta', + sessionId: 's1', messageId: 'm1', blockIndex: 0, text: 'hello', @@ -1815,6 +1943,7 @@ describe('StreamStateTracker', () => { it('returns empty additional array for tool_use messages', () => { const result = tracker.processMessage({ type: 'tool_use', + sessionId: 's1', toolName: 'read_file', toolInput: {}, toolUseId: 'tu-1', @@ -1825,6 +1954,7 @@ describe('StreamStateTracker', () => { it('returns empty additional array for error messages', () => { const result = tracker.processMessage({ type: 'error', + sessionId: 's1', message: 'err', errorType: DroidErrorType.ERROR, timestamp: 't', @@ -1840,17 +1970,20 @@ describe('StreamStateTracker', () => { tracker1.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); const r2 = tracker2.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(r2.additional).toEqual([]); const r1 = tracker1.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(r1.additional).toHaveLength(1); @@ -1860,10 +1993,12 @@ describe('StreamStateTracker', () => { it('simulates multi-turn session with reset between turns', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); tracker.processMessage({ type: 'token_usage_update', + sessionId: 's1', inputTokens: 50, outputTokens: 25, cacheReadTokens: 0, @@ -1872,6 +2007,7 @@ describe('StreamStateTracker', () => { }); const turn1Result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(turn1Result.additional).toHaveLength(1); @@ -1884,10 +2020,12 @@ describe('StreamStateTracker', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.ExecutingTool, }); tracker.processMessage({ type: 'token_usage_update', + sessionId: 's1', inputTokens: 200, outputTokens: 100, cacheReadTokens: 10, @@ -1896,6 +2034,7 @@ describe('StreamStateTracker', () => { }); const turn2Result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(turn2Result.additional).toHaveLength(1); @@ -1910,16 +2049,19 @@ describe('StreamStateTracker', () => { it('handles rapid idle→non-idle→idle transitions', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(result.additional).toHaveLength(1); @@ -1929,10 +2071,12 @@ describe('StreamStateTracker', () => { it('handles WaitingForToolConfirmation as non-idle', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.WaitingForToolConfirmation, }); const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(result.additional).toHaveLength(1); @@ -1942,10 +2086,12 @@ describe('StreamStateTracker', () => { it('handles CompactingConversation as non-idle', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.CompactingConversation, }); const result = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(result.additional).toHaveLength(1); @@ -1955,17 +2101,20 @@ describe('StreamStateTracker', () => { it('multiple idle transitions after non-idle only emit first Result (no duplicate)', () => { tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.StreamingAssistantMessage, }); const r1 = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(r1.additional).toHaveLength(1); const r2 = tracker.processMessage({ type: 'working_state_changed', + sessionId: 's1', state: DroidWorkingState.Idle, }); expect(r2.additional).toHaveLength(0);