Skip to content

Keep user-opened tools expanded across streaming updates #743

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
draxxris wants to merge 1 commit into
base: main
Choose a base branch
from
+42 −10
Open

Keep user-opened tools expanded across streaming updates #743

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 17 additions & 1 deletion components/ChatWindow.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Set<string>>(() => 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,
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -1193,7 +1209,7 @@ export function ChatWindow({ session, searchTarget, onSearchTargetHandled, initi
);
})()}
{streamState.isStreaming && hasStreamingContent && streamState.streamingMessage && (
<MessageView message={streamState.streamingMessage as AgentMessage} isStreaming modelNames={modelNames} cwd={messageCwd} onOpenFile={onOpenFile} onOpenSession={onOpenSession} />
<MessageView message={streamState.streamingMessage as AgentMessage} isStreaming modelNames={modelNames} cwd={messageCwd} onOpenFile={onOpenFile} onOpenSession={onOpenSession} expandedToolIds={expandedToolIds} onToggleTool={handleToggleTool} />
)}

{agentRunning && !hasStreamingContent && agentPhase && (
Expand Down
34 changes: 25 additions & 9 deletions components/MessageView.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<string>;
onToggleTool?: (toolCallId: string) => void;
}

function formatTime(ts?: number): string | null {
Expand Down Expand Up @@ -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 <UserMessageView message={message as UserMessage} cwd={cwd} onOpenFile={onOpenFile} entryId={entryId} onFork={onFork} forking={forking} onNavigate={onNavigate} prevAssistantEntryId={prevAssistantEntryId} onEditContent={onEditContent} />;
}
if (message.role === "assistant") {
return <AssistantMessageView message={message as AssistantMessage} isStreaming={isStreaming} toolResults={toolResults} modelNames={modelNames} cwd={cwd} onOpenFile={onOpenFile} onOpenSession={onOpenSession} showTimestamp={showTimestamp} prevTimestamp={prevTimestamp} sessionId={sessionId} entryId={entryId} searchBlock={searchBlock} writtenFiles={writtenFiles} />;
return <AssistantMessageView message={message as AssistantMessage} isStreaming={isStreaming} toolResults={toolResults} modelNames={modelNames} cwd={cwd} onOpenFile={onOpenFile} onOpenSession={onOpenSession} showTimestamp={showTimestamp} prevTimestamp={prevTimestamp} sessionId={sessionId} entryId={entryId} searchBlock={searchBlock} writtenFiles={writtenFiles} expandedToolIds={expandedToolIds} onToggleTool={onToggleTool} />;
}
if (message.role === "toolResult") {
// Rendered inline under its toolCall — skip standalone rendering if paired
Expand Down Expand Up @@ -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 }: {
Expand Down Expand Up @@ -589,6 +594,8 @@ function AssistantMessageView({
entryId,
searchBlock,
writtenFiles,
expandedToolIds,
onToggleTool,
}: {
message: AssistantMessage;
isStreaming?: boolean;
Expand All @@ -603,6 +610,8 @@ function AssistantMessageView({
entryId?: string;
searchBlock?: AssistantContentBlock;
writtenFiles?: WrittenFile[];
expandedToolIds?: Set<string>;
onToggleTool?: (toolCallId: string) => void;
}) {
const { t } = useI18n();
const time = showTimestamp ? formatTime(message.timestamp) : null;
Expand Down Expand Up @@ -782,7 +791,7 @@ function AssistantMessageView({

<div style={{ display: "flex", flexDirection: "column", gap: 8 }}>
{blockItems.map(({ block, originalIndex }) => (
<BlockView key={`${entryId ?? "stream"}-${originalIndex}`} block={block} searchTarget={block === searchBlock} toolResults={toolResults} isStreaming={isStreaming} streamingDuration={streamingDurations.get(originalIndex) ?? (block.type === "thinking" ? thinkingDurationFromFile : undefined)} toolCallDurations={toolCallDurations} cwd={cwd} onOpenFile={onOpenFile} onOpenSession={onOpenSession} sessionId={sessionId} entryId={entryId} blockIndex={originalIndex} />
<BlockView key={`${entryId ?? "stream"}-${originalIndex}`} block={block} searchTarget={block === searchBlock} toolResults={toolResults} isStreaming={isStreaming} streamingDuration={streamingDurations.get(originalIndex) ?? (block.type === "thinking" ? thinkingDurationFromFile : undefined)} toolCallDurations={toolCallDurations} cwd={cwd} onOpenFile={onOpenFile} onOpenSession={onOpenSession} sessionId={sessionId} entryId={entryId} blockIndex={originalIndex} expandedToolIds={expandedToolIds} onToggleTool={onToggleTool} />
))}
</div>

Expand Down Expand Up @@ -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<string, ToolResultMessage>; isStreaming?: boolean; streamingDuration?: number; toolCallDurations?: Map<string, number>; 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<string, ToolResultMessage>; isStreaming?: boolean; streamingDuration?: number; toolCallDurations?: Map<string, number>; cwd?: string; onOpenFile?: (filePath: string) => void; onOpenSession?: (sessionId: string) => void; sessionId?: string; entryId?: string; blockIndex: number; expandedToolIds?: Set<string>; onToggleTool?: (toolCallId: string) => void }) {
if (block.type === "text") {
return <div data-message-text data-search-target={searchTarget || undefined}><TextBlock block={block as TextContent} isStreaming={isStreaming} cwd={cwd} onOpenFile={onOpenFile} /></div>;
}
Expand All @@ -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 <ToolCallBlock block={tc} result={result} duration={duration} onOpenSession={onOpenSession} />;
const fallbackKey = `${entryId ?? "stream"}-${blockIndex}`;
const toggleId = tc.toolCallId || fallbackKey;
const isExpanded = expandedToolIds ? expandedToolIds.has(toggleId) : undefined;
const handleToggle = onToggleTool ? () => onToggleTool(toggleId) : undefined;
return <ToolCallBlock block={tc} result={result} duration={duration} onOpenSession={onOpenSession} expanded={isExpanded} onToggle={handleToggle} />;
}
return null;
}
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -1031,7 +1047,7 @@ function ToolCallBlock({ block, result, duration, onOpenSession }: { block: Tool
{/* ── Tool call header ── */}
<div style={{ display: "flex", alignItems: "stretch", minWidth: 0 }}>
<button
onClick={() => setExpanded((v) => !v)}
onClick={handleToggle}
style={{
display: "flex",
alignItems: "center",
Expand Down