diff --git a/src/web-ui/src/flow_chat/services/EventBatcher.test.ts b/src/web-ui/src/flow_chat/services/EventBatcher.test.ts index 8db9fe3bf..1d4029afa 100644 --- a/src/web-ui/src/flow_chat/services/EventBatcher.test.ts +++ b/src/web-ui/src/flow_chat/services/EventBatcher.test.ts @@ -1,6 +1,12 @@ import { afterEach, describe, expect, it } from 'vitest'; import { setIncludeSensitiveDiagnostics } from '@/shared/utils/logger'; -import { getBatchedEventsLogPayload, summarizeBatchedEventsForLog, type BatchedEvent } from './EventBatcher'; +import { + generateToolEventKey, + getBatchedEventsLogPayload, + summarizeBatchedEventsForLog, + type BatchedEvent, + type ToolEventData, +} from './EventBatcher'; describe('summarizeBatchedEventsForLog', () => { afterEach(() => { @@ -65,3 +71,24 @@ describe('summarizeBatchedEventsForLog', () => { expect(summaryText).not.toContain('src/secret.ts'); }); }); + +describe('generateToolEventKey', () => { + it('accumulates Write params so argument deltas survive batching', () => { + const keyInfo = generateToolEventKey({ + sessionId: 'session-1', + turnId: 'turn-1', + roundId: 'round-1', + toolEvent: { + event_type: 'ParamsPartial', + tool_id: 'tool-1', + tool_name: 'Write', + params: '{"file_path":"src/app.ts"', + }, + } satisfies ToolEventData); + + expect(keyInfo).toEqual({ + key: 'tool:params:session-1:tool-1', + strategy: 'accumulate', + }); + }); +}); diff --git a/src/web-ui/src/flow_chat/services/EventBatcher.ts b/src/web-ui/src/flow_chat/services/EventBatcher.ts index d706b36a4..a33ad64e9 100644 --- a/src/web-ui/src/flow_chat/services/EventBatcher.ts +++ b/src/web-ui/src/flow_chat/services/EventBatcher.ts @@ -396,11 +396,9 @@ export function generateToolEventKey(data: ToolEventData): { key: string; strate } if (eventType === 'ParamsPartial') { - const toolName = (toolEvent as any).tool_name || ''; - const isWriteLike = ['write', 'write_notebook', 'file_write', 'Write'].includes(toolName); return { key: `tool:params:${sessionId}:${toolUseId}`, - strategy: isWriteLike ? 'replace' : 'accumulate' + strategy: 'accumulate' }; } if (eventType === 'Progress') { diff --git a/src/web-ui/src/flow_chat/services/FlowChatManager.test.ts b/src/web-ui/src/flow_chat/services/FlowChatManager.test.ts index ef342f24b..d3d62257d 100644 --- a/src/web-ui/src/flow_chat/services/FlowChatManager.test.ts +++ b/src/web-ui/src/flow_chat/services/FlowChatManager.test.ts @@ -4,6 +4,10 @@ import { FlowChatManager } from './FlowChatManager'; const storeMocks = vi.hoisted(() => ({ store: {} as any, initializeEventListeners: vi.fn(), + eventBatchers: [] as Array<{ + flushNow: ReturnType; + destroy: ReturnType; + }>, })); vi.mock('./ProcessingStatusManager', () => ({ @@ -32,7 +36,12 @@ vi.mock('../state-machine', () => ({ vi.mock('./EventBatcher', () => ({ EventBatcher: class { - constructor(private readonly options: { onFlush: (events: Array<{ key: string; payload: unknown }>) => void }) {} + public flushNow = vi.fn(); + public destroy = vi.fn(); + + constructor(private readonly options: { onFlush: (events: Array<{ key: string; payload: unknown }>) => void }) { + storeMocks.eventBatchers.push(this); + } flush(events: Array<{ key: string; payload: unknown }>): void { this.options.onFlush(events); @@ -104,9 +113,70 @@ describe('FlowChatManager initialization', () => { beforeEach(() => { (FlowChatManager as any).instance = undefined; vi.clearAllMocks(); + storeMocks.eventBatchers.length = 0; storeMocks.initializeEventListeners.mockResolvedValue(() => {}); }); + it('flushes and destroys the batcher when the singleton is disposed', () => { + storeMocks.store = {}; + + const manager = FlowChatManager.getInstance(); + const batcher = storeMocks.eventBatchers[0]; + + FlowChatManager.disposeInstance(); + + expect(batcher.flushNow).toHaveBeenCalledTimes(1); + expect(batcher.destroy).toHaveBeenCalledTimes(1); + expect(batcher.flushNow.mock.invocationCallOrder[0]).toBeLessThan( + batcher.destroy.mock.invocationCallOrder[0], + ); + }); + + it('runs listener cleanup if disposal wins the initialization race', async () => { + storeMocks.store = {}; + const listenerInitialization = createDeferred<() => void>(); + const cleanup = vi.fn(); + storeMocks.initializeEventListeners.mockReturnValue(listenerInitialization.promise); + + const manager = FlowChatManager.getInstance(); + const initializeListeners = (manager as any).initializeEventListeners(); + + await flushAsyncWork(); + manager.destroy(); + expect(cleanup).not.toHaveBeenCalled(); + + listenerInitialization.resolve(cleanup); + await initializeListeners; + + expect(cleanup).toHaveBeenCalledTimes(1); + expect((manager as any).eventListenerInitialized).toBe(false); + expect((manager as any).eventListenerCleanup).toBeNull(); + }); + + it('stops workspace initialization if the manager is disposed while listeners initialize', async () => { + const listenerInitialization = createDeferred<() => void>(); + storeMocks.initializeEventListeners.mockReturnValue(listenerInitialization.promise); + storeMocks.store = { + registerPersistUnreadCompletionCallback: vi.fn(), + loadSessionMetadataPage: vi.fn(async () => ({ + sessions: [], + totalTopLevelCount: 0, + hasMore: false, + })), + }; + + const manager = FlowChatManager.getInstance(); + const initialize = manager.initialize('D:/workspace/BitFun'); + + await flushAsyncWork(); + manager.destroy(); + listenerInitialization.resolve(vi.fn()); + + await expect(initialize).resolves.toBe(false); + expect(storeMocks.store.registerPersistUnreadCompletionCallback).not.toHaveBeenCalled(); + expect(storeMocks.store.loadSessionMetadataPage).not.toHaveBeenCalled(); + }); + it('reuses concurrent initialization for the same workspace history restore', async () => { const metadataLoad = createDeferred<{ sessions: unknown[]; diff --git a/src/web-ui/src/flow_chat/services/FlowChatManager.ts b/src/web-ui/src/flow_chat/services/FlowChatManager.ts index 68fb06068..43d6b7a94 100644 --- a/src/web-ui/src/flow_chat/services/FlowChatManager.ts +++ b/src/web-ui/src/flow_chat/services/FlowChatManager.ts @@ -48,7 +48,7 @@ import { const log = createLogger('FlowChatManager'); export class FlowChatManager { - private static instance: FlowChatManager; + private static instance: FlowChatManager | null = null; private context: FlowChatContext; private agentService: AgentService; private eventListenerInitialized = false; @@ -56,6 +56,7 @@ export class FlowChatManager { private eventListenerCleanup: (() => void) | null = null; private initializationRequests = new Map>(); private latestInitializationRequestKey: string | null = null; + private disposed = false; private constructor() { this.context = { @@ -96,12 +97,22 @@ export class FlowChatManager { return FlowChatManager.instance; } + public static disposeInstance(): void { + FlowChatManager.instance?.destroy(); + FlowChatManager.instance = null; + } + async initialize( workspacePath: string, preferredMode?: string, remoteConnectionId?: string, remoteSshHost?: string ): Promise { + if (this.disposed) { + log.debug('Ignoring initialize call on disposed FlowChatManager', { workspacePath }); + return false; + } + const requestKey = FlowChatManager.createInitializationRequestKey( workspacePath, preferredMode, @@ -153,6 +164,9 @@ export class FlowChatManager { ): Promise { try { await this.initializeEventListeners(); + if (this.disposed) { + return false; + } // Register callback to persist unread completion changes to backend this.context.flowChatStore.registerPersistUnreadCompletionCallback( @@ -171,6 +185,9 @@ export class FlowChatManager { remoteSshHost, 'flow_chat_manager' ); + if (this.disposed) { + return false; + } const sessionMatchesWorkspace = (session: { workspacePath?: string; @@ -207,6 +224,9 @@ export class FlowChatManager { remoteSshHost, 'flow_chat_manager_preferred_mode' ); + if (this.disposed) { + return false; + } state = this.context.flowChatStore.getState(); workspaceSessions = Array.from(state.sessions.values()).filter(sessionMatchesWorkspace); if (workspaceSessions.some(session => session.mode === preferredMode) || !nextPage.hasMore) { @@ -251,6 +271,9 @@ export class FlowChatManager { latestSession.remoteSshHost, { deferFullHistoryUntilActive: true }, ); + if (this.disposed) { + return false; + } } if (!isCurrentInitializationRequest()) { @@ -289,6 +312,9 @@ export class FlowChatManager { } private async initializeEventListeners(): Promise { + if (this.disposed) { + return; + } if (this.eventListenerInitialized) { return; } @@ -297,11 +323,17 @@ export class FlowChatManager { } this.eventListenerInitializationPromise = (async () => { - this.eventListenerCleanup = await initializeEventListeners( + const cleanup = await initializeEventListeners( this.context, (sessionId, turnId, result) => this.handleTodoWriteResult(sessionId, turnId, result) ); + if (this.disposed) { + cleanup(); + return; + } + + this.eventListenerCleanup = cleanup; this.eventListenerInitialized = true; })(); @@ -321,7 +353,24 @@ export class FlowChatManager { this.eventListenerInitializationPromise = null; } + public destroy(): void { + if (this.disposed) { + return; + } + + this.context.eventBatcher.flushNow(); + this.disposed = true; + this.initializationRequests.clear(); + this.latestInitializationRequestKey = null; + this.cleanupEventListeners(); + this.context.eventBatcher.destroy(); + } + private processBatchedEvents(events: Array<{ key: string; payload: any }>): void { + if (this.disposed) { + return; + } + processBatchedEvents( this.context, events, @@ -680,5 +729,12 @@ export class FlowChatManager { } } } + +if (import.meta.hot) { + import.meta.hot.dispose(() => { + FlowChatManager.disposeInstance(); + }); +} + export const flowChatManager = FlowChatManager.getInstance(); export default flowChatManager; diff --git a/src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.test.ts b/src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.test.ts index 8ac37bbf1..01f3bd120 100644 --- a/src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.test.ts +++ b/src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.test.ts @@ -62,6 +62,38 @@ describe('resolveDialogTurnDisplayContent', () => { }); }); +describe('mergeParamsPartialEventData', () => { + it('appends Write argument deltas within a batch', () => { + const merged = __test_only__.mergeParamsPartialEventData( + { + sessionId: 'session-1', + turnId: 'turn-1', + roundId: 'round-1', + toolEvent: { + event_type: 'ParamsPartial', + tool_id: 'tool-1', + tool_name: 'Write', + params: '{"file_path":"src/app.ts","content":"', + }, + }, + { + sessionId: 'session-1', + turnId: 'turn-1', + roundId: 'round-1', + toolEvent: { + event_type: 'ParamsPartial', + tool_id: 'tool-1', + tool_name: 'Write', + params: 'hello', + }, + }, + ); + + expect((merged.toolEvent as any).params).toBe('{"file_path":"src/app.ts","content":"hello'); + }); + +}); + describe('shouldProcessEvent', () => { const mockSessionId = 'test-session'; const mockTurnId = 'test-turn'; diff --git a/src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.ts b/src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.ts index b2f957c78..26dcf8faa 100644 --- a/src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.ts +++ b/src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.ts @@ -123,8 +123,29 @@ function resolveDialogTurnDisplayContent( return resolveThreadGoalUserMessageDisplay(base, metadata); } +function mergeParamsPartialEventData( + existing: ToolEventData, + incoming: ToolEventData, +): ToolEventData { + const existingToolEvent = existing.toolEvent as ParamsPartialToolEvent; + const incomingToolEvent = incoming.toolEvent as ParamsPartialToolEvent; + const existingParams = normalizeParamsPartialFragment(existingToolEvent.params); + const incomingParams = normalizeParamsPartialFragment(incomingToolEvent.params); + + return { + ...existing, + ...incoming, + toolEvent: { + ...existingToolEvent, + ...incomingToolEvent, + params: existingParams + incomingParams, + }, + }; +} + export const __test_only__ = { resolveDialogTurnDisplayContent, + mergeParamsPartialEventData, }; function shouldMarkUnreadCompletion(sessionId: string): boolean { @@ -1619,15 +1640,7 @@ function handleToolEvent( key, eventData, 'accumulate', - (existing, incoming) => ({ - ...existing, - toolEvent: { - ...(existing.toolEvent as ParamsPartialToolEvent), - params: - normalizeParamsPartialFragment((existing.toolEvent as ParamsPartialToolEvent).params) + - normalizeParamsPartialFragment((incoming.toolEvent as ParamsPartialToolEvent).params) - } - }) + mergeParamsPartialEventData, ); } else { context.eventBatcher.add(key, eventData, 'replace'); diff --git a/src/web-ui/src/flow_chat/services/flow-chat-manager/ToolEventModule.ts b/src/web-ui/src/flow_chat/services/flow-chat-manager/ToolEventModule.ts index 85fb370e5..0ebb83010 100644 --- a/src/web-ui/src/flow_chat/services/flow-chat-manager/ToolEventModule.ts +++ b/src/web-ui/src/flow_chat/services/flow-chat-manager/ToolEventModule.ts @@ -199,8 +199,7 @@ function applyParamsPartial( if (!incomingParams) { return; } - const isWriteFullParamsSnapshot = isWriteTool && incomingParams.trimStart().startsWith('{'); - const newBuffer = isWriteFullParamsSnapshot ? incomingParams : prevBuffer + incomingParams; + const newBuffer = prevBuffer + incomingParams; let parsedParams: Record = {}; try { diff --git a/src/web-ui/src/flow_chat/tool-cards/FileOperationToolCard.test.tsx b/src/web-ui/src/flow_chat/tool-cards/FileOperationToolCard.test.tsx index 5f4a170e4..cb4631df8 100644 --- a/src/web-ui/src/flow_chat/tool-cards/FileOperationToolCard.test.tsx +++ b/src/web-ui/src/flow_chat/tool-cards/FileOperationToolCard.test.tsx @@ -19,6 +19,7 @@ const mocks = vi.hoisted(() => ({ createDiffEditorTab: vi.fn(), openFile: vi.fn(), codePreviewProps: [] as Array>, + inlineDiffPreviewProps: [] as Array>, getOperationDiff: vi.fn(async () => ({ originalContent: '', modifiedContent: '', @@ -79,7 +80,10 @@ vi.mock('../components/CodePreview', () => ({ })); vi.mock('../components/InlineDiffPreview', () => ({ - InlineDiffPreview: ({ modifiedContent }: { modifiedContent: string }) => 
{modifiedContent}
, + InlineDiffPreview: (props: Record) => { + mocks.inlineDiffPreviewProps.push(props); + return 
{String(props.modifiedContent ?? '')}
; + }, })); vi.mock('../../shared/utils/tabUtils', () => ({ @@ -140,6 +144,7 @@ describe('FileOperationToolCard', () => { mocks.createDiffEditorTab.mockReset(); mocks.openFile.mockReset(); mocks.codePreviewProps = []; + mocks.inlineDiffPreviewProps = []; mocks.useGitState.mockClear(); mocks.useGitState.mockReturnValue({ isRepository: false, @@ -739,4 +744,124 @@ describe('FileOperationToolCard', () => { autoScrollToBottom: false, }); }); + + it('keeps completed write preview compact while auto-collapsing from streaming', async () => { + const config: ToolCardConfig = { + toolName: 'Write', + displayName: 'Write', + icon: 'WRITE', + requiresConfirmation: false, + resultDisplayType: 'detailed', + description: 'Write a file', + displayMode: 'standard', + }; + const streamingToolItem: FlowToolItem = { + id: 'tool-1', + type: 'tool', + toolName: 'Write', + status: 'streaming', + isParamsStreaming: true, + toolCall: { + id: 'call-1', + name: 'Write', + input: { + file_path: 'src/generated.ts', + content: 'line 1\nline 2\nline 3\nline 4\nline 5\nline 6', + }, + }, + partialParams: { + file_path: 'src/generated.ts', + content: 'line 1\nline 2\nline 3\nline 4\nline 5\nline 6', + }, + } as FlowToolItem; + const completedToolItem: FlowToolItem = { + ...streamingToolItem, + status: 'completed', + isParamsStreaming: false, + toolResult: { + success: true, + result: { + file_path: 'src/generated.ts', + }, + }, + } as FlowToolItem; + + await act(async () => { + root.render( + + ); + }); + + mocks.inlineDiffPreviewProps = []; + + await act(async () => { + root.render( + + ); + }); + + expect(mocks.inlineDiffPreviewProps.length).toBeGreaterThan(0); + expect(mocks.inlineDiffPreviewProps.map(props => props.maxHeight)).not.toContain(330); + expect(mocks.inlineDiffPreviewProps.map(props => props.maxHeight)).toContain(88); + }); + + it('uses the larger diff preview height after a completed write card is manually expanded', async () => { + const toolItem: FlowToolItem = { + id: 'tool-1', + type: 'tool', + toolName: 'Write', + status: 'completed', + isParamsStreaming: false, + toolCall: { + id: 'call-1', + name: 'Write', + input: { + file_path: 'src/generated.ts', + content: 'line 1\nline 2\nline 3\nline 4\nline 5\nline 6', + }, + }, + toolResult: { + success: true, + result: { + file_path: 'src/generated.ts', + }, + }, + } as FlowToolItem; + const config: ToolCardConfig = { + toolName: 'Write', + displayName: 'Write', + icon: 'WRITE', + requiresConfirmation: false, + resultDisplayType: 'detailed', + description: 'Write a file', + displayMode: 'standard', + }; + + await act(async () => { + root.render( + + ); + }); + + mocks.inlineDiffPreviewProps = []; + + const card = container.querySelector('.base-tool-card') as HTMLDivElement | null; + await act(async () => { + card?.dispatchEvent(new dom.window.MouseEvent('click', { bubbles: true })); + }); + + expect(mocks.inlineDiffPreviewProps.map(props => props.maxHeight)).toContain(330); + }); }); diff --git a/src/web-ui/src/flow_chat/tool-cards/FileOperationToolCard.tsx b/src/web-ui/src/flow_chat/tool-cards/FileOperationToolCard.tsx index 71db928cd..d86144faa 100644 --- a/src/web-ui/src/flow_chat/tool-cards/FileOperationToolCard.tsx +++ b/src/web-ui/src/flow_chat/tool-cards/FileOperationToolCard.tsx @@ -442,6 +442,10 @@ export const FileOperationToolCard: React.FC = ({ }, [sessionId, toolCall?.id, status, isFailed]); const isLoading = status === 'preparing' || status === 'streaming' || status === 'running'; + const shouldUseExpandedDiffPreviewHeight = + status === 'completed' && + isContentExpanded && + previousStatusRef.current === status; const previewVariant = useMemo(() => { if (toolItem.toolName === 'Edit') { if (status !== 'completed' && newStringContent) { @@ -835,7 +839,7 @@ export const FileOperationToolCard: React.FC = ({ const renderExpandedContent = () => { if (isFailed) return null; - const previewMaxHeight = status === 'completed' + const previewMaxHeight = shouldUseExpandedDiffPreviewHeight ? FILE_OPERATION_DIFF_MAX_HEIGHT : FILE_OPERATION_STREAMING_MAX_HEIGHT;