diff --git a/packages/desktop/src/agent/MessageStream.tsx b/packages/desktop/src/agent/MessageStream.tsx index fe93bbe..48a7427 100644 --- a/packages/desktop/src/agent/MessageStream.tsx +++ b/packages/desktop/src/agent/MessageStream.tsx @@ -1,4 +1,5 @@ import { useEffect, useRef, useState, useCallback, useMemo } from 'react'; +import { Copy, Check } from 'lucide-react'; import { useVirtualizer } from '@tanstack/react-virtual'; import { useGlobalStore } from '../stores/global.store'; import MessageItem from '../shared/MessageItem'; @@ -6,6 +7,7 @@ import UnifiedDiffView from '../shared/UnifiedDiffView'; import type { Item } from '@shared/types'; import type { CheckpointDiff } from '../lib/core-api'; import { useAgentApproval, useAgentRollback } from '../hooks/useAgent'; +import { useCopyToClipboard } from '../hooks/useCopyToClipboard'; interface MessageStreamProps { threadId: string; @@ -232,6 +234,7 @@ export default function MessageStream({ threadId }: MessageStreamProps) { forkThread, revertedFilesByTurnId, } = useAgentRollback(); + const { copiedId, copy } = useCopyToClipboard(); const parentRef = useRef(null); const markScopeRestored = useGlobalStore((s) => s.markScopeRestored); const markFileRestored = useGlobalStore((s) => s.markFileRestored); @@ -253,7 +256,7 @@ export default function MessageStream({ threadId }: MessageStreamProps) { [turns] ); - const { renderEntries, callIdToToolName, entryCountByTurnId, turnById } = useMemo(() => { + const { renderEntries, callIdToToolName, entryCountByTurnId, turnById, assistantContentByTurnId } = useMemo(() => { const entries: Array<{ item: Item; turnId: string; @@ -264,16 +267,23 @@ export default function MessageStream({ threadId }: MessageStreamProps) { const nameMap: Record = {}; const countMap = new Map(); const turnMap = new Map(); + const contentMap = new Map(); for (const turn of turns) { turnMap.set(turn.id, turn); + const assistantParts: string[] = []; for (const item of turn.items) { if (item.type === 'tool_result') { toolResultByCallId[item.callId] = item as any; } else if (item.type === 'tool_call') { nameMap[item.id] = item.name; + } else if (item.type === 'message' && item.role === 'assistant' && item.content) { + assistantParts.push(item.content); } } + if (assistantParts.length > 0) { + contentMap.set(turn.id, assistantParts.join('\n\n')); + } } for (const turn of turns) { for (const item of turn.items) { @@ -297,6 +307,7 @@ export default function MessageStream({ threadId }: MessageStreamProps) { callIdToToolName: nameMap, entryCountByTurnId: countMap, turnById: turnMap, + assistantContentByTurnId: contentMap, }; }, [turnsStructureKey]); @@ -526,6 +537,35 @@ export default function MessageStream({ threadId }: MessageStreamProps) { toolResult={entry.toolResult} /> + {isLastInTurn && (() => { + const assistantContent = assistantContentByTurnId.get(entry.turnId); + const isTurnDone = turn?.status === 'completed' || turn?.status === 'error'; + if (!assistantContent || !isTurnDone) return null; + const isCopied = copiedId === `turn-${entry.turnId}`; + return ( +
+
+ +
+
+ ); + })()} {isLastInTurn && (
-
- {isAssistant && messageContent != null && } - {item.partial && ( - - )} -
- {isAssistant && !item.partial && messageContent != null && ( -
- + {messageContent != null && ( +
+ + {item.partial && ( + + )}
)}
diff --git a/packages/desktop/test/assistant-content-by-turn.test.ts b/packages/desktop/test/assistant-content-by-turn.test.ts new file mode 100644 index 0000000..486c446 --- /dev/null +++ b/packages/desktop/test/assistant-content-by-turn.test.ts @@ -0,0 +1,144 @@ +import { describe, it, expect } from 'vitest'; +import type { Item, Turn } from '../shared/types'; + +function buildAssistantContentByTurnId(turns: Turn[]): Map { + const contentMap = new Map(); + for (const turn of turns) { + const assistantParts: string[] = []; + for (const item of turn.items) { + if (item.type === 'message' && item.role === 'assistant' && item.content) { + assistantParts.push(item.content); + } + } + if (assistantParts.length > 0) { + contentMap.set(turn.id, assistantParts.join('\n\n')); + } + } + return contentMap; +} + +function makeMsg(role: 'user' | 'assistant', content: string): Item { + return { id: 'm-' + content, type: 'message', role, content }; +} + +function makeToolCall( + name: string, + status: 'pending' | 'running' | 'approved' | 'rejected', + id?: string +): Item { + return { id: id ?? 'tc-' + name, type: 'tool_call', name, args: {}, status }; +} + +function makeToolResult(callId: string, output?: string): Item { + return { id: 'tr-' + callId, type: 'tool_result', callId, output: output ?? 'ok' }; +} + +describe('assistantContentByTurnId', () => { + it('concatenates multiple assistant messages in one turn with double newline', () => { + const turns: Turn[] = [ + { + id: 't1', + items: [ + makeMsg('user', 'hello'), + makeMsg('assistant', 'I will help'), + makeToolCall('read_file', 'approved', 'tc-1'), + makeToolResult('tc-1', 'file content'), + makeMsg('assistant', 'Here is the result'), + ], + status: 'completed', + }, + ]; + const map = buildAssistantContentByTurnId(turns); + expect(map.get('t1')).toBe('I will help\n\nHere is the result'); + }); + + it('returns single message content when only one assistant message exists', () => { + const turns: Turn[] = [ + { + id: 't1', + items: [makeMsg('user', 'hi'), makeMsg('assistant', 'hello')], + status: 'completed', + }, + ]; + const map = buildAssistantContentByTurnId(turns); + expect(map.get('t1')).toBe('hello'); + }); + + it('does not include turn with no assistant messages', () => { + const turns: Turn[] = [ + { + id: 't1', + items: [makeMsg('user', 'hi')], + status: 'completed', + }, + ]; + const map = buildAssistantContentByTurnId(turns); + expect(map.has('t1')).toBe(false); + }); + + it('handles multiple turns independently', () => { + const turns: Turn[] = [ + { + id: 't1', + items: [makeMsg('user', 'a'), makeMsg('assistant', 'reply-a')], + status: 'completed', + }, + { + id: 't2', + items: [makeMsg('user', 'b'), makeMsg('assistant', 'reply-b1'), makeMsg('assistant', 'reply-b2')], + status: 'completed', + }, + ]; + const map = buildAssistantContentByTurnId(turns); + expect(map.get('t1')).toBe('reply-a'); + expect(map.get('t2')).toBe('reply-b1\n\nreply-b2'); + }); + + it('ignores empty assistant messages', () => { + const turns: Turn[] = [ + { + id: 't1', + items: [ + makeMsg('user', 'hi'), + { id: 'm-empty', type: 'message', role: 'assistant', content: '' }, + makeMsg('assistant', 'actual content'), + ], + status: 'completed', + }, + ]; + const map = buildAssistantContentByTurnId(turns); + expect(map.get('t1')).toBe('actual content'); + }); + + it('handles interrupted turn with error status', () => { + const turns: Turn[] = [ + { + id: 't1', + items: [ + makeMsg('user', 'do something'), + makeMsg('assistant', 'partial response'), + makeToolCall('bash', 'running', 'tc-1'), + ], + status: 'error', + }, + ]; + const map = buildAssistantContentByTurnId(turns); + expect(map.get('t1')).toBe('partial response'); + }); + + it('handles turn with only tool calls and no assistant text', () => { + const turns: Turn[] = [ + { + id: 't1', + items: [ + makeMsg('user', 'run it'), + makeToolCall('bash', 'approved', 'tc-1'), + makeToolResult('tc-1', 'done'), + ], + status: 'completed', + }, + ]; + const map = buildAssistantContentByTurnId(turns); + expect(map.has('t1')).toBe(false); + }); +});