From 459e8183fd41d3aae53ca9bd48b767dea6846ff3 Mon Sep 17 00:00:00 2001 From: draxxris <24984408+draxxris@users.noreply.github.com> Date: Wed, 2 Sep 2026 06:59:59 -0400 Subject: [PATCH] fix(chat): keep expanded tool display open across streaming updates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lift expanded state from ToolCallBlock's local useState to a chat-level Set keyed by toolCallId (fallback to entry/stream index). The local state was tied to the component instance and reset on every streaming delta and when a streamed tool was promoted to history via loadSession, causing an open tool to collapse on each new LLM line. A session-scoped map in ChatWindow now survives re-renders and the stream → history transition, and is cleared on session change. --- components/ChatWindow.tsx | 18 +++++++++++++++++- components/MessageView.tsx | 34 +++++++++++++++++++++++++--------- 2 files changed, 42 insertions(+), 10 deletions(-) diff --git a/components/ChatWindow.tsx b/components/ChatWindow.tsx index 466aa3730..39a7fea6d 100644 --- a/components/ChatWindow.tsx +++ b/components/ChatWindow.tsx @@ -270,6 +270,20 @@ export function ChatWindow({ session, searchTarget, onSearchTargetHandled, initi return position && !position.atBottom ? position : null; }); const [restoreAnchorReady, setRestoreAnchorReady] = useState(false); + // Lifted expanded state for tool calls — survives streaming re-renders and + // stream → history promotion (keyed by toolCallId, fallback to stream/entry index). + const [expandedToolIds, setExpandedToolIds] = useState>(() => new Set()); + const handleToggleTool = useCallback((id: string) => { + setExpandedToolIds((prev) => { + const next = new Set(prev); + if (next.has(id)) next.delete(id); + else next.add(id); + return next; + }); + }, []); + useEffect(() => { + setExpandedToolIds(new Set()); + }, [session?.id]); const { loading, error, messages, entryIds, historyCursor, hasEarlierMessages, streamState, @@ -1060,6 +1074,8 @@ export function ChatWindow({ session, searchTarget, onSearchTargetHandled, initi prevTimestamp={idx > 0 ? (messages[idx - 1] as AgentMessage & { timestamp?: number }).timestamp : undefined} sessionId={session?.id ?? sessionIdRef.current ?? undefined} writtenFiles={options.writtenFiles} + expandedToolIds={expandedToolIds} + onToggleTool={handleToggleTool} /> ); if (!isVisible || currentRefIdx === undefined) return view; @@ -1193,7 +1209,7 @@ export function ChatWindow({ session, searchTarget, onSearchTargetHandled, initi ); })()} {streamState.isStreaming && hasStreamingContent && streamState.streamingMessage && ( - + )} {agentRunning && !hasStreamingContent && agentPhase && ( diff --git a/components/MessageView.tsx b/components/MessageView.tsx index 354c67097..69a212b63 100644 --- a/components/MessageView.tsx +++ b/components/MessageView.tsx @@ -204,6 +204,9 @@ interface Props { * final answer text-only. */ writtenFiles?: WrittenFile[]; + /** Lifted expanded state for tool calls — when provided, ToolCallBlock becomes controlled. */ + expandedToolIds?: Set; + onToggleTool?: (toolCallId: string) => void; } function formatTime(ts?: number): string | null { @@ -252,12 +255,12 @@ function haveSameRelevantToolResults( return true; } -export const MessageView = memo(function MessageView({ message, isStreaming, toolResults, modelNames, cwd, onOpenFile, onOpenSession, entryId, searchBlock, onFork, forking, onNavigate, prevAssistantEntryId, onEditContent, showTimestamp, prevTimestamp, sessionId, writtenFiles }: Props) { +export const MessageView = memo(function MessageView({ message, isStreaming, toolResults, modelNames, cwd, onOpenFile, onOpenSession, entryId, searchBlock, onFork, forking, onNavigate, prevAssistantEntryId, onEditContent, showTimestamp, prevTimestamp, sessionId, writtenFiles, expandedToolIds, onToggleTool }: Props) { if (message.role === "user") { return ; } if (message.role === "assistant") { - return ; + return ; } if (message.role === "toolResult") { // Rendered inline under its toolCall — skip standalone rendering if paired @@ -291,7 +294,9 @@ export const MessageView = memo(function MessageView({ message, isStreaming, too && prev.showTimestamp === next.showTimestamp && prev.prevTimestamp === next.prevTimestamp && prev.writtenFiles === next.writtenFiles - && prev.sessionId === next.sessionId; + && prev.sessionId === next.sessionId + && prev.expandedToolIds === next.expandedToolIds + && prev.onToggleTool === next.onToggleTool; }); function UserMessageView({ message, cwd, onOpenFile, entryId, onFork, forking, onNavigate, prevAssistantEntryId, onEditContent }: { @@ -589,6 +594,8 @@ function AssistantMessageView({ entryId, searchBlock, writtenFiles, + expandedToolIds, + onToggleTool, }: { message: AssistantMessage; isStreaming?: boolean; @@ -603,6 +610,8 @@ function AssistantMessageView({ entryId?: string; searchBlock?: AssistantContentBlock; writtenFiles?: WrittenFile[]; + expandedToolIds?: Set; + onToggleTool?: (toolCallId: string) => void; }) { const { t } = useI18n(); const time = showTimestamp ? formatTime(message.timestamp) : null; @@ -782,7 +791,7 @@ function AssistantMessageView({
{blockItems.map(({ block, originalIndex }) => ( - + ))}
@@ -860,7 +869,7 @@ function AssistantMessageView({ ); } -function BlockView({ block, searchTarget, toolResults, isStreaming, streamingDuration, toolCallDurations, cwd, onOpenFile, onOpenSession, sessionId, entryId, blockIndex }: { block: AssistantContentBlock; searchTarget?: boolean; toolResults?: Map; isStreaming?: boolean; streamingDuration?: number; toolCallDurations?: Map; cwd?: string; onOpenFile?: (filePath: string) => void; onOpenSession?: (sessionId: string) => void; sessionId?: string; entryId?: string; blockIndex: number }) { +function BlockView({ block, searchTarget, toolResults, isStreaming, streamingDuration, toolCallDurations, cwd, onOpenFile, onOpenSession, sessionId, entryId, blockIndex, expandedToolIds, onToggleTool }: { block: AssistantContentBlock; searchTarget?: boolean; toolResults?: Map; isStreaming?: boolean; streamingDuration?: number; toolCallDurations?: Map; cwd?: string; onOpenFile?: (filePath: string) => void; onOpenSession?: (sessionId: string) => void; sessionId?: string; entryId?: string; blockIndex: number; expandedToolIds?: Set; onToggleTool?: (toolCallId: string) => void }) { if (block.type === "text") { return
; } @@ -871,7 +880,11 @@ function BlockView({ block, searchTarget, toolResults, isStreaming, streamingDur const tc = block as ToolCallContent; const result = toolResults?.get(tc.toolCallId); const duration = toolCallDurations?.get(tc.toolCallId); - return ; + const fallbackKey = `${entryId ?? "stream"}-${blockIndex}`; + const toggleId = tc.toolCallId || fallbackKey; + const isExpanded = expandedToolIds ? expandedToolIds.has(toggleId) : undefined; + const handleToggle = onToggleTool ? () => onToggleTool(toggleId) : undefined; + return ; } return null; } @@ -1001,9 +1014,12 @@ function isSubagentToolDetails(value: unknown): value is SubagentToolDetails { return details.kind === "pi-web-subagent" && typeof details.sessionId === "string"; } -function ToolCallBlock({ block, result, duration, onOpenSession }: { block: ToolCallContent; result?: ToolResultMessage; duration?: number; onOpenSession?: (sessionId: string) => void }) { +function ToolCallBlock({ block, result, duration, onOpenSession, expanded: controlledExpanded, onToggle }: { block: ToolCallContent; result?: ToolResultMessage; duration?: number; onOpenSession?: (sessionId: string) => void; expanded?: boolean; onToggle?: () => void }) { const { t } = useI18n(); - const [expanded, setExpanded] = useState(false); + const [localExpanded, setLocalExpanded] = useState(false); + const isControlled = controlledExpanded !== undefined && onToggle !== undefined; + const expanded = isControlled ? controlledExpanded : localExpanded; + const handleToggle = isControlled ? onToggle! : () => setLocalExpanded((v) => !v); const inputStr = getToolCallInputText(block); const isStreamingInput = block.rawInput !== undefined; const isEditTool = isEditToolName(block.toolName); @@ -1031,7 +1047,7 @@ function ToolCallBlock({ block, result, duration, onOpenSession }: { block: Tool {/* ── Tool call header ── */}