feat(ui): add thinking ui to mothership (#4254)

* feat(ui): Add thinking ui

* fix tests

* Remove duplicate helper for block timing

* fix lint

* fix endedAt timestamp bug

* fix stuck subagent thinking
This commit is contained in:
Theodore Li
2026-04-22 16:52:56 -07:00
committed by GitHub
parent 8ce56fe1f2
commit f7ab39984c
16 changed files with 520 additions and 41 deletions

View File

@@ -52,6 +52,8 @@ const ContentBlockSchema = z.object({
lifecycle: z.enum(['start', 'end']).optional(),
status: z.enum(['complete', 'error', 'cancelled']).optional(),
toolCall: StoredToolCallSchema.optional(),
timestamp: z.number().optional(),
endedAt: z.number().optional(),
})
const StopSchema = z.object({

View File

@@ -5,10 +5,12 @@ import { ChevronDown, Expandable, ExpandableContent, PillsRing } from '@/compone
import { cn } from '@/lib/core/utils/cn'
import type { ToolCallData } from '../../../../types'
import { getAgentIcon } from '../../utils'
import { ThinkingBlock } from '../thinking-block'
import { ToolCallItem } from './tool-call-item'
export type AgentGroupItem =
| { type: 'text'; content: string }
| { type: 'thinking'; content: string; startedAt?: number; endedAt?: number }
| { type: 'tool'; data: ToolCallData }
interface AgentGroupProps {
@@ -16,6 +18,7 @@ interface AgentGroupProps {
agentLabel: string
items: AgentGroupItem[]
isDelegating?: boolean
isStreaming?: boolean
autoCollapse?: boolean
defaultExpanded?: boolean
}
@@ -35,6 +38,7 @@ export function AgentGroup({
agentLabel,
items,
isDelegating = false,
isStreaming = false,
autoCollapse = false,
defaultExpanded = false,
}: AgentGroupProps) {
@@ -110,16 +114,39 @@ export function AgentGroup({
<Expandable expanded={expanded}>
<ExpandableContent>
<div className='flex flex-col gap-1.5 pt-0.5'>
{items.map((item, idx) =>
item.type === 'tool' ? (
<ToolCallItem
key={item.data.id}
toolName={item.data.toolName}
displayTitle={item.data.displayTitle}
status={item.data.status}
streamingArgs={item.data.streamingArgs}
/>
) : (
{items.map((item, idx) => {
if (item.type === 'tool') {
return (
<ToolCallItem
key={item.data.id}
toolName={item.data.toolName}
displayTitle={item.data.displayTitle}
status={item.data.status}
streamingArgs={item.data.streamingArgs}
/>
)
}
if (item.type === 'thinking') {
const elapsedMs =
item.startedAt !== undefined && item.endedAt !== undefined
? item.endedAt - item.startedAt
: undefined
if (elapsedMs !== undefined && elapsedMs <= 3000) return null
return (
<div key={`thinking-${idx}`} className='pl-6'>
<ThinkingBlock
content={item.content}
isActive={
isStreaming && idx === items.length - 1 && item.endedAt === undefined
}
isStreaming={isStreaming}
startedAt={item.startedAt}
endedAt={item.endedAt}
/>
</div>
)
}
return (
<span
key={`text-${idx}`}
className='pl-6 font-base text-[var(--text-secondary)] text-small'
@@ -127,7 +154,7 @@ export function AgentGroup({
{item.content.trim()}
</span>
)
)}
})}
</div>
</ExpandableContent>
</Expandable>

View File

@@ -3,3 +3,4 @@ export { AgentGroup, CircleStop } from './agent-group'
export { ChatContent } from './chat-content'
export { Options } from './options'
export { PendingTagIndicator, parseSpecialTags, SpecialTags } from './special-tags'
export { ThinkingBlock } from './thinking-block'

View File

@@ -0,0 +1 @@
export { ThinkingBlock } from './thinking-block'

View File

@@ -0,0 +1,123 @@
'use client'
import { useEffect, useLayoutEffect, useRef, useState } from 'react'
import { ChevronDown, Expandable, ExpandableContent } from '@/components/emcn'
import { BrainIcon } from '@/components/icons'
import { cn } from '@/lib/core/utils/cn'
interface ThinkingBlockProps {
content: string
isActive: boolean
isStreaming?: boolean
startedAt?: number
endedAt?: number
}
const MIN_VISIBLE_THINKING_MS = 3000
export function ThinkingBlock({
content,
isActive,
isStreaming = false,
startedAt,
endedAt,
}: ThinkingBlockProps) {
// Start collapsed so the `Expandable` plays its height-open animation
// when `expanded` flips to true below — otherwise the panel mounts
// already-open and jumps up with its full content in one frame.
const [expanded, setExpanded] = useState(false)
const panelRef = useRef<HTMLDivElement>(null)
const wasActiveRef = useRef<boolean | null>(null)
// Suppress active thinking until it exceeds MIN_VISIBLE_THINKING_MS.
// Completed-<=threshold is filtered upstream in message-content, so if
// we're mounted with isActive=false we've already passed that gate.
const [thresholdReached, setThresholdReached] = useState(() => {
if (!isActive || startedAt === undefined) return true
return Date.now() - startedAt > MIN_VISIBLE_THINKING_MS
})
useEffect(() => {
if (thresholdReached) return
if (!isActive || startedAt === undefined) {
setThresholdReached(true)
return
}
const remainingMs = Math.max(0, MIN_VISIBLE_THINKING_MS - (Date.now() - startedAt))
const id = window.setTimeout(() => setThresholdReached(true), remainingMs + 50)
return () => window.clearTimeout(id)
}, [isActive, startedAt, thresholdReached])
useEffect(() => {
// Wait until the threshold has actually been reached — otherwise this
// effect fires during the 3-second hidden period (while the component
// returns null) and sets `expanded` to true before the panel is even
// rendered, so the Collapsible mounts already-open with no animation.
if (!thresholdReached) return
if (wasActiveRef.current === isActive) return
// On first run (wasActiveRef === null): open if the stream is live —
// even when thinking itself has already ended — so a mid-stream refresh
// shows the thinking panel open while the rest of the response is still
// being generated. Subsequent runs only react to the isActive transition
// (auto-collapse when thinking ends).
const isFirstRun = wasActiveRef.current === null
wasActiveRef.current = isActive
const target = isFirstRun ? isActive || isStreaming : isActive
// Defer to the next frame so Radix Collapsible paints the closed state
// first, then sees the transition to open. Without this, React can batch
// the mount + flip into a single commit and the animation never plays.
const id = window.requestAnimationFrame(() => setExpanded(target))
return () => window.cancelAnimationFrame(id)
}, [isActive, isStreaming, thresholdReached])
useLayoutEffect(() => {
if (!isActive || !expanded) return
const el = panelRef.current
if (!el) return
el.scrollTop = el.scrollHeight
}, [content, isActive, expanded])
if (!thresholdReached) return null
const elapsedMs =
startedAt !== undefined && endedAt !== undefined && endedAt >= startedAt
? endedAt - startedAt
: undefined
const elapsedSeconds =
elapsedMs !== undefined ? Math.max(1, Math.round(elapsedMs / 1000)) : undefined
const label = isActive
? 'Thinking'
: elapsedSeconds !== undefined
? `Thought for ${elapsedSeconds}s`
: 'Thought'
return (
<div className='flex flex-col gap-1.5'>
<button
type='button'
onClick={() => setExpanded((prev) => !prev)}
className='flex cursor-pointer items-center gap-2'
>
<div className='flex h-[16px] w-[16px] flex-shrink-0 items-center justify-center'>
<BrainIcon className='h-[14px] w-[14px] text-[var(--text-icon)]' />
</div>
<span className='font-base text-[var(--text-body)] text-sm'>{label}</span>
<ChevronDown
className={cn(
'h-[7px] w-[9px] text-[var(--text-icon)] transition-transform duration-150',
!expanded && '-rotate-90'
)}
/>
</button>
<Expandable expanded={expanded}>
<ExpandableContent>
<div ref={panelRef} className='max-h-[110px] overflow-y-scroll pt-0.5 pr-2 pl-6'>
<div className='whitespace-pre-wrap break-words font-base text-[13px] text-[var(--text-secondary)] leading-[18px] opacity-60'>
{content}
</div>
</div>
</ExpandableContent>
</Expandable>
</div>
)
}

View File

@@ -10,7 +10,14 @@ import { ClientToolCallState } from '@/lib/copilot/tools/client/tool-call-state'
import type { ContentBlock, MothershipResource, OptionItem, ToolCallData } from '../../types'
import { SUBAGENT_LABELS, TOOL_UI_METADATA } from '../../types'
import type { AgentGroupItem } from './components'
import { AgentGroup, ChatContent, CircleStop, Options, PendingTagIndicator } from './components'
import {
AgentGroup,
ChatContent,
CircleStop,
Options,
PendingTagIndicator,
ThinkingBlock,
} from './components'
const FILE_SUBAGENT_ID = 'file'
@@ -19,6 +26,14 @@ interface TextSegment {
content: string
}
interface ThinkingSegment {
type: 'thinking'
id: string
content: string
startedAt?: number
endedAt?: number
}
interface AgentGroupSegment {
type: 'agent_group'
id: string
@@ -38,7 +53,12 @@ interface StoppedSegment {
type: 'stopped'
}
type MessageSegment = TextSegment | AgentGroupSegment | OptionsSegment | StoppedSegment
type MessageSegment =
| TextSegment
| ThinkingSegment
| AgentGroupSegment
| OptionsSegment
| StoppedSegment
const SUBAGENT_KEYS = new Set(Object.keys(SUBAGENT_LABELS))
@@ -156,6 +176,46 @@ function parseBlocks(blocks: ContentBlock[]): MessageSegment[] {
continue
}
if (block.type === 'subagent_thinking') {
if (!block.content || !group) continue
group.isDelegating = false
const lastItem = group.items[group.items.length - 1]
if (lastItem?.type === 'thinking' && lastItem.endedAt === undefined) {
lastItem.content += block.content
if (block.endedAt !== undefined) lastItem.endedAt = block.endedAt
} else {
group.items.push({
type: 'thinking',
content: block.content,
startedAt: block.timestamp,
endedAt: block.endedAt,
})
}
continue
}
if (block.type === 'thinking') {
if (!block.content?.trim()) continue
if (group) {
pushGroup(group)
group = null
}
const last = segments[segments.length - 1]
if (last?.type === 'thinking' && last.endedAt === undefined) {
last.content += block.content
if (block.endedAt !== undefined) last.endedAt = block.endedAt
} else {
segments.push({
type: 'thinking',
id: `thinking-${i}`,
content: block.content,
startedAt: block.timestamp,
endedAt: block.endedAt,
})
}
continue
}
if (block.type === 'text') {
if (!block.content) continue
if (block.subagent) {
@@ -383,7 +443,9 @@ export function MessageContent({
const hasSubagentEnded = blocks.some((b) => b.type === 'subagent_end')
const showTrailingThinking =
isStreaming && !hasTrailingContent && (hasSubagentEnded || allLastGroupToolsDone)
isStreaming &&
!hasTrailingContent &&
(lastSegment.type === 'thinking' || hasSubagentEnded || allLastGroupToolsDone)
const lastOpenSubagentGroupId = [...segments]
.reverse()
.find(
@@ -405,6 +467,30 @@ export function MessageContent({
onWorkspaceResourceSelect={onWorkspaceResourceSelect}
/>
)
case 'thinking': {
const isActive =
isStreaming && i === segments.length - 1 && segment.endedAt === undefined
const elapsedMs =
segment.startedAt !== undefined && segment.endedAt !== undefined
? segment.endedAt - segment.startedAt
: undefined
// Hide completed thinking that took 3s or less — quick thinking
// isn't worth the visual noise. Still show while active (unknown
// duration yet) and still show when timing is missing (old
// persisted blocks) so we don't drop historical content.
if (elapsedMs !== undefined && elapsedMs <= 3000) return null
return (
<div key={segment.id} className={isStreaming ? 'animate-stream-fade-in' : undefined}>
<ThinkingBlock
content={segment.content}
isActive={isActive}
isStreaming={isStreaming}
startedAt={segment.startedAt}
endedAt={segment.endedAt}
/>
</div>
)
}
case 'agent_group': {
const toolItems = segment.items.filter((item) => item.type === 'tool')
const allToolsDone =
@@ -419,6 +505,7 @@ export function MessageContent({
agentLabel={segment.agentLabel}
items={segment.items}
isDelegating={segment.isDelegating}
isStreaming={isStreaming}
autoCollapse={allToolsDone && hasFollowingText}
defaultExpanded={segment.id === lastOpenSubagentGroupId}
/>

View File

@@ -11,7 +11,7 @@ import type {
PersistedFileAttachment,
PersistedMessage,
} from '@/lib/copilot/chat/persisted-message'
import { normalizeMessage } from '@/lib/copilot/chat/persisted-message'
import { normalizeMessage, withBlockTiming } from '@/lib/copilot/chat/persisted-message'
import { resolveStreamToolOutcome } from '@/lib/copilot/chat/stream-tool-outcome'
import { MOTHERSHIP_CHAT_API_PATH, STREAM_STORAGE_KEY } from '@/lib/copilot/constants'
import type {
@@ -26,6 +26,7 @@ import {
MothershipStreamV1SessionKind,
MothershipStreamV1SpanLifecycleEvent,
MothershipStreamV1SpanPayloadKind,
MothershipStreamV1TextChannel,
MothershipStreamV1ToolOutcome,
MothershipStreamV1ToolPhase,
MothershipStreamV1ToolStatus,
@@ -699,11 +700,37 @@ function parseStreamBatchResponse(value: unknown): StreamBatchResponse {
}
function toRawPersistedContentBlock(block: ContentBlock): Record<string, unknown> | null {
const persisted = toRawPersistedContentBlockBody(block)
return persisted ? withBlockTiming(persisted, block) : null
}
function toRawPersistedContentBlockBody(block: ContentBlock): Record<string, unknown> | null {
switch (block.type) {
case 'text':
return {
type: MothershipStreamV1EventType.text,
...(block.subagent ? { lane: 'subagent' } : {}),
channel: MothershipStreamV1TextChannel.assistant,
content: block.content ?? '',
}
case 'thinking':
return {
type: MothershipStreamV1EventType.text,
channel: MothershipStreamV1TextChannel.thinking,
content: block.content ?? '',
}
case 'subagent_thinking':
return {
type: MothershipStreamV1EventType.text,
lane: 'subagent',
channel: MothershipStreamV1TextChannel.thinking,
content: block.content ?? '',
}
case 'subagent_text':
return {
type: MothershipStreamV1EventType.text,
lane: 'subagent',
channel: MothershipStreamV1TextChannel.assistant,
content: block.content ?? '',
}
case 'tool_call':
@@ -773,22 +800,27 @@ function buildAssistantSnapshotMessage(params: {
}
function markMessageStopped(message: PersistedMessage): PersistedMessage {
if (!message.contentBlocks?.some((block) => block.toolCall?.state === 'executing')) {
const hasExecutingTool = message.contentBlocks?.some(
(block) => block.toolCall?.state === 'executing'
)
const hasOpenBlock = message.contentBlocks?.some((block) => block.endedAt === undefined)
if (!hasExecutingTool && !hasOpenBlock) {
return message
}
const nextBlocks = message.contentBlocks.map((block) => {
if (block.toolCall?.state !== 'executing') {
return block
const stopTs = Date.now()
const nextBlocks = (message.contentBlocks ?? []).map((block) => {
const stamped = block.endedAt === undefined ? { ...block, endedAt: stopTs } : block
if (stamped.toolCall?.state !== 'executing') {
return stamped
}
return {
...block,
...stamped,
toolCall: {
...block.toolCall,
...stamped.toolCall,
state: 'cancelled' as const,
display: {
...(block.toolCall.display ?? {}),
...(stamped.toolCall.display ?? {}),
title: 'Stopped by user',
},
},
@@ -1716,10 +1748,34 @@ export function useChat(
streamingBlocksRef.current = []
}
const ensureTextBlock = (subagentName?: string): ContentBlock => {
const toEventMs = (ts: string | undefined): number => {
if (ts) {
const parsed = Date.parse(ts)
if (Number.isFinite(parsed)) return parsed
}
return Date.now()
}
const stampBlockEnd = (block: ContentBlock | undefined, ts?: string) => {
if (block && block.endedAt === undefined) block.endedAt = toEventMs(ts)
}
const ensureTextBlock = (subagentName: string | undefined, ts?: string): ContentBlock => {
const last = blocks[blocks.length - 1]
if (last?.type === 'text' && last.subagent === subagentName) return last
const b: ContentBlock = { type: 'text', content: '' }
stampBlockEnd(last, ts)
const b: ContentBlock = { type: 'text', content: '', timestamp: toEventMs(ts) }
if (subagentName) b.subagent = subagentName
blocks.push(b)
return b
}
const ensureThinkingBlock = (subagentName: string | undefined, ts?: string): ContentBlock => {
const targetType = subagentName ? 'subagent_thinking' : 'thinking'
const last = blocks[blocks.length - 1]
if (last?.type === targetType && last.subagent === subagentName) return last
stampBlockEnd(last, ts)
const b: ContentBlock = { type: targetType, content: '', timestamp: toEventMs(ts) }
if (subagentName) b.subagent = subagentName
blocks.push(b)
return b
@@ -1737,9 +1793,9 @@ export function useChat(
return activeSubagent
}
const appendInlineErrorTag = (tag: string, subagentName?: string) => {
const appendInlineErrorTag = (tag: string, subagentName?: string, ts?: string) => {
if (runningText.includes(tag)) return
const tb = ensureTextBlock(subagentName)
const tb = ensureTextBlock(subagentName, ts)
const prefix = runningText.length > 0 && !runningText.endsWith('\n') ? '\n' : ''
tb.content = `${tb.content ?? ''}${prefix}${tag}`
runningText += `${prefix}${tag}`
@@ -1950,13 +2006,20 @@ export function useChat(
case MothershipStreamV1EventType.text: {
const chunk = parsed.payload.text
if (chunk) {
const eventTs = typeof parsed.ts === 'string' ? parsed.ts : undefined
if (parsed.payload.channel === MothershipStreamV1TextChannel.thinking) {
const tb = ensureThinkingBlock(scopedSubagent, eventTs)
tb.content = (tb.content ?? '') + chunk
flushText()
break
}
const contentSource: 'main' | 'subagent' = scopedSubagent ? 'subagent' : 'main'
const needsBoundaryNewline =
lastContentSource !== null &&
lastContentSource !== contentSource &&
runningText.length > 0 &&
!runningText.endsWith('\n')
const tb = ensureTextBlock(scopedSubagent)
const tb = ensureTextBlock(scopedSubagent, eventTs)
const normalizedChunk = needsBoundaryNewline ? `\n${chunk}` : chunk
tb.content = (tb.content ?? '') + normalizedChunk
runningText += normalizedChunk
@@ -2170,6 +2233,7 @@ export function useChat(
output: payload.output,
error: typeof payload.error === 'string' ? payload.error : undefined,
}
stampBlockEnd(blocks[idx])
flush()
if (tc.name === ReadTool.id && tc.status === 'success') {
@@ -2292,6 +2356,7 @@ export function useChat(
}
if (!toolMap.has(id)) {
stampBlockEnd(blocks[blocks.length - 1])
toolMap.set(id, blocks.length)
blocks.push({
type: 'tool_call',
@@ -2303,6 +2368,7 @@ export function useChat(
params: args,
calledBy: scopedSubagent,
},
timestamp: Date.now(),
})
if (name === ReadTool.id || isResourceToolName(name)) {
if (args) toolArgsMap.set(id, args)
@@ -2376,6 +2442,7 @@ export function useChat(
if (payload.kind === MothershipStreamV1RunKind.compaction_start) {
const compactionId = `compaction_${Date.now()}`
activeCompactionId = compactionId
stampBlockEnd(blocks[blocks.length - 1])
toolMap.set(compactionId, blocks.length)
blocks.push({
type: 'tool_call',
@@ -2385,6 +2452,7 @@ export function useChat(
status: 'executing',
displayTitle: 'Compacting context...',
},
timestamp: Date.now(),
})
flush()
} else if (payload.kind === MothershipStreamV1RunKind.compaction_done) {
@@ -2394,8 +2462,10 @@ export function useChat(
if (idx !== undefined && blocks[idx]?.toolCall) {
blocks[idx].toolCall!.status = 'success'
blocks[idx].toolCall!.displayTitle = 'Compacted context'
stampBlockEnd(blocks[idx])
} else {
toolMap.set(compactionId, blocks.length)
const endNow = Date.now()
blocks.push({
type: 'tool_call',
toolCall: {
@@ -2404,6 +2474,8 @@ export function useChat(
status: 'success',
displayTitle: 'Compacted context',
},
timestamp: endNow,
endedAt: endNow,
})
}
flush()
@@ -2432,7 +2504,8 @@ export function useChat(
activeSubagent = name
activeSubagentParentToolCallId = parentToolCallId
if (!isSameActiveSubagent) {
blocks.push({ type: 'subagent', content: name })
stampBlockEnd(blocks[blocks.length - 1])
blocks.push({ type: 'subagent', content: name, timestamp: Date.now() })
}
if (name === FILE_SUBAGENT_ID && !isSameActiveSubagent) {
applyPreviewSessionUpdate({
@@ -2472,7 +2545,18 @@ export function useChat(
activeSubagent = undefined
activeSubagentParentToolCallId = undefined
}
blocks.push({ type: 'subagent_end' })
const endNow = Date.now()
if (name) {
for (let i = blocks.length - 1; i >= 0; i--) {
const b = blocks[i]
if (b.type === 'subagent' && b.content === name && b.endedAt === undefined) {
b.endedAt = endNow
break
}
}
}
stampBlockEnd(blocks[blocks.length - 1])
blocks.push({ type: 'subagent_end', timestamp: endNow })
flush()
}
break
@@ -2480,11 +2564,16 @@ export function useChat(
case MothershipStreamV1EventType.error: {
sawStreamError = true
setError(parsed.payload.message || parsed.payload.error || 'An error occurred')
appendInlineErrorTag(buildInlineErrorTag(parsed.payload), scopedSubagent)
appendInlineErrorTag(
buildInlineErrorTag(parsed.payload),
scopedSubagent,
typeof parsed.ts === 'string' ? parsed.ts : undefined
)
break
}
case MothershipStreamV1EventType.complete: {
sawCompleteEvent = true
stampBlockEnd(blocks[blocks.length - 1])
// `complete` is the end-of-turn marker; drain whatever
// else arrived in the same TCP chunk (trailing text,
// followups, run metadata) before stopping. Do NOT
@@ -2888,6 +2977,10 @@ export function useChat(
const sourceBlocks = overrides?.blocks ?? streamingBlocksRef.current
const storedBlocks = sourceBlocks.map((block) => {
const timing = {
...(typeof block.timestamp === 'number' ? { timestamp: block.timestamp } : {}),
...(typeof block.endedAt === 'number' ? { endedAt: block.endedAt } : {}),
}
if (block.type === 'tool_call' && block.toolCall) {
const isCancelled =
block.toolCall.status === 'executing' || block.toolCall.status === 'cancelled'
@@ -2905,9 +2998,10 @@ export function useChat(
...(display ? { display } : {}),
calledBy: block.toolCall.calledBy,
},
...timing,
}
}
return { type: block.type, content: block.content }
return { type: block.type, content: block.content, ...timing }
})
if (storedBlocks.length > 0) {
@@ -3465,11 +3559,21 @@ export function useChat(
queryClient.getQueryData<TaskChatHistory>(taskKeys.detail(chatIdRef.current))
?.activeStreamId ||
undefined
// Snapshot the active assistant message id BEFORE clearActiveTurn()
// nulls the ref. Used below to restrict markMessageStopped to the
// in-flight turn only — historical messages from the chat history
// also lack `endedAt` on their legacy blocks (pre-timing-fields),
// and without this gate we'd corrupt them with cancelled markers.
const activeAssistantMessageId =
activeTurnRef.current?.assistantMessageId ??
(sid ? getLiveAssistantMessageId(sid) : undefined)
const stopContentSnapshot = streamingContentRef.current
const stopNow = Date.now()
const stopBlocksSnapshot = streamingBlocksRef.current.map((block) => ({
...block,
...(block.options ? { options: [...block.options] } : {}),
...(block.toolCall ? { toolCall: { ...block.toolCall } } : {}),
...(block.endedAt === undefined ? { endedAt: stopNow } : {}),
}))
// Snapshot BEFORE clearActiveTurn() nulls the refs. Both
// persistPartialResponse and the abort/stop fetches run inside
@@ -3491,22 +3595,31 @@ export function useChat(
await queryClient.cancelQueries({ queryKey: taskKeys.detail(activeChatId) })
upsertTaskChatHistory(activeChatId, (current) => ({
...current,
messages: current.messages.map(markMessageStopped),
messages: current.messages.map((message) =>
activeAssistantMessageId && message.id === activeAssistantMessageId
? markMessageStopped(message)
: message
),
}))
} else {
setPendingMessages((prev) =>
prev.map((msg) => {
if (!msg.contentBlocks?.some((block) => block.toolCall?.status === 'executing')) {
const hasExecutingTool = msg.contentBlocks?.some(
(block) => block.toolCall?.status === 'executing'
)
const hasOpenBlock = msg.contentBlocks?.some((block) => block.endedAt === undefined)
if (!hasExecutingTool && !hasOpenBlock) {
return msg
}
const updatedBlocks = msg.contentBlocks.map((block) => {
if (block.toolCall?.status !== 'executing') {
return block
const updatedBlocks = (msg.contentBlocks ?? []).map((block) => {
const stamped = block.endedAt === undefined ? { ...block, endedAt: stopNow } : block
if (stamped.toolCall?.status !== 'executing') {
return stamped
}
return {
...block,
...stamped,
toolCall: {
...block.toolCall,
...stamped.toolCall,
status: 'cancelled' as const,
displayTitle: 'Stopped by user',
},

View File

@@ -114,6 +114,7 @@ export interface OptionItem {
export const ContentBlockType = {
text: 'text',
thinking: 'thinking',
tool_call: 'tool_call',
subagent: 'subagent',
subagent_end: 'subagent_end',
@@ -130,6 +131,8 @@ export interface ContentBlock {
subagent?: string
toolCall?: ToolCallInfo
options?: OptionItem[]
timestamp?: number
endedAt?: number
}
export interface ChatMessageAttachment {

View File

@@ -102,10 +102,27 @@ export function useAutoScroll(
rafIdRef.current = requestAnimationFrame(guardedScroll)
}
// CSS-driven height animations (e.g. Radix Collapsible expanding
// mid-stream) grow scrollHeight without triggering MutationObserver,
// so auto-scroll stops following. When any animation starts in the
// container, follow rAF for a short window so the container stays
// pinned to the bottom while the animation runs.
const onAnimationStart = () => {
if (!stickyRef.current) return
const until = performance.now() + 500
const follow = () => {
if (performance.now() > until || !stickyRef.current) return
scrollToBottom()
requestAnimationFrame(follow)
}
requestAnimationFrame(follow)
}
el.addEventListener('wheel', onWheel, { passive: true })
el.addEventListener('touchstart', onTouchStart, { passive: true })
el.addEventListener('touchmove', onTouchMove, { passive: true })
el.addEventListener('scroll', onScroll, { passive: true })
el.addEventListener('animationstart', onAnimationStart)
const observer = new MutationObserver(onMutation)
observer.observe(el, { childList: true, subtree: true, characterData: true })
@@ -115,6 +132,7 @@ export function useAutoScroll(
el.removeEventListener('touchstart', onTouchStart)
el.removeEventListener('touchmove', onTouchMove)
el.removeEventListener('scroll', onScroll)
el.removeEventListener('animationstart', onAnimationStart)
observer.disconnect()
cancelAnimationFrame(rafIdRef.current)
if (stickyRef.current) scrollToBottom()

View File

@@ -16,6 +16,7 @@ import {
ToolCallStatus,
} from '@/app/workspace/[workspaceId]/home/types'
import type { PersistedContentBlock, PersistedMessage } from './persisted-message'
import { withBlockTiming } from './persisted-message'
const STATE_TO_STATUS: Record<string, ToolCallStatus> = {
[MothershipStreamV1ToolOutcome.success]: ToolCallStatus.success,
@@ -44,6 +45,11 @@ function toToolCallInfo(block: PersistedContentBlock): ToolCallInfo | undefined
}
function toDisplayBlock(block: PersistedContentBlock): ContentBlock | undefined {
const displayed = toDisplayBlockBody(block)
return displayed ? withBlockTiming(displayed, block) : undefined
}
function toDisplayBlockBody(block: PersistedContentBlock): ContentBlock | undefined {
switch (block.type) {
case MothershipStreamV1EventType.text:
if (block.lane === 'subagent') {
@@ -52,6 +58,9 @@ function toDisplayBlock(block: PersistedContentBlock): ContentBlock | undefined
}
return { type: ContentBlockType.subagent_text, content: block.content }
}
if (block.channel === 'thinking') {
return { type: ContentBlockType.thinking, content: block.content }
}
return { type: ContentBlockType.text, content: block.content }
case MothershipStreamV1EventType.tool:
if (!toToolCallInfo(block)) return undefined

View File

@@ -12,6 +12,7 @@ import {
describe('persisted-message', () => {
it('round-trips canonical tool blocks through normalizeMessage', () => {
const blockTimestamp = 1_700_000_000_000
const result: OrchestratorResult = {
success: true,
content: 'done',
@@ -19,7 +20,7 @@ describe('persisted-message', () => {
contentBlocks: [
{
type: 'tool_call',
timestamp: Date.now(),
timestamp: blockTimestamp,
calledBy: 'workflow',
toolCall: {
id: 'tool-1',
@@ -41,6 +42,7 @@ describe('persisted-message', () => {
{
type: 'tool',
phase: 'call',
timestamp: blockTimestamp,
toolCall: {
id: 'tool-1',
name: 'read',

View File

@@ -38,6 +38,8 @@ export interface PersistedContentBlock {
status?: MothershipStreamV1CompletionStatus
content?: string
toolCall?: PersistedToolCall
timestamp?: number
endedAt?: number
}
export interface PersistedFileAttachment {
@@ -85,7 +87,25 @@ function resolveToolState(block: ContentBlock): PersistedToolState {
return tc.status as PersistedToolState
}
/**
* Copy `timestamp` / `endedAt` from a source object onto a target object.
* Shared by every block mapper (persist, display, snapshot) so the timing
* metadata that drives the `Thought for Ns` chip survives the full
* persist → normalize → display round-trip — and one rule lives in one place.
*/
export function withBlockTiming<T>(target: T, src: { timestamp?: number; endedAt?: number }): T {
const writable = target as { timestamp?: number; endedAt?: number }
if (typeof src.timestamp === 'number') writable.timestamp = src.timestamp
if (typeof src.endedAt === 'number') writable.endedAt = src.endedAt
return target
}
function mapContentBlock(block: ContentBlock): PersistedContentBlock {
const persisted = mapContentBlockBody(block)
return withBlockTiming(persisted, block)
}
function mapContentBlockBody(block: ContentBlock): PersistedContentBlock {
switch (block.type) {
case 'text':
return {
@@ -242,6 +262,8 @@ interface RawBlock {
kind?: string
lifecycle?: string
status?: string
timestamp?: number
endedAt?: number
toolCall?: {
id?: string
name?: string
@@ -406,7 +428,16 @@ function normalizeLegacyBlock(block: RawBlock): PersistedContentBlock {
}
function normalizeBlock(block: RawBlock): PersistedContentBlock {
return isCanonicalBlock(block) ? normalizeCanonicalBlock(block) : normalizeLegacyBlock(block)
const result = isCanonicalBlock(block)
? normalizeCanonicalBlock(block)
: normalizeLegacyBlock(block)
if (typeof block.timestamp === 'number' && result.timestamp === undefined) {
result.timestamp = block.timestamp
}
if (typeof block.endedAt === 'number' && result.endedAt === undefined) {
result.endedAt = block.endedAt
}
return result
}
function normalizeLegacyToolCall(tc: LegacyToolCall): PersistedContentBlock {

View File

@@ -24,6 +24,10 @@ import {
sseHandlers,
subAgentHandlers,
} from '@/lib/copilot/request/handlers'
import {
flushSubagentThinkingBlock,
flushThinkingBlock,
} from '@/lib/copilot/request/handlers/types'
import { getCopilotTracer } from '@/lib/copilot/request/otel'
import {
eventToStreamEvent,
@@ -337,6 +341,13 @@ export async function runStreamLoop(
const subagentName = streamEvent.payload.agent
const spanEvt = streamEvent.payload.event
const isPendingPause = spanData?.pending === true
// A subagent lifecycle boundary breaks the main thinking stream.
// Flush any open thinking block into contentBlocks BEFORE we push
// the `subagent` marker, or the persisted order ends up
// [subagent, thinking] and the UI renders the subagent group
// above a thinking block that actually happened first.
flushSubagentThinkingBlock(context)
flushThinkingBlock(context)
if (spanEvt === MothershipStreamV1SpanLifecycleEvent.start) {
const lastParent = context.subAgentParentStack[context.subAgentParentStack.length - 1]
const lastBlock = context.contentBlocks[context.contentBlocks.length - 1]
@@ -377,6 +388,19 @@ export async function runStreamLoop(
context.subAgentParentStack.length > 0
? context.subAgentParentStack[context.subAgentParentStack.length - 1]
: undefined
if (subagentName) {
for (let i = context.contentBlocks.length - 1; i >= 0; i--) {
const b = context.contentBlocks[i]
if (
b.type === 'subagent' &&
b.content === subagentName &&
b.endedAt === undefined
) {
b.endedAt = Date.now()
break
}
}
}
return
}
}
@@ -434,6 +458,12 @@ export async function runStreamLoop(
endedOn = CopilotSseCloseReason.Aborted
}
}
// An abort or error can tear down the loop mid-thinking. Flush any
// open thinking blocks so partial-persistence on /chat/stop sees
// them in contentBlocks with endedAt stamped, instead of silently
// dropping the in-flight reasoning.
flushSubagentThinkingBlock(context)
flushThinkingBlock(context)
clearTimeout(timeoutId)
// Legacy TraceCollector span (consumed by the in-memory trace

View File

@@ -36,6 +36,8 @@ import {
addContentBlock,
emitSyntheticToolResult,
ensureTerminalToolCallState,
flushSubagentThinkingBlock,
flushThinkingBlock,
getScopedParentToolCallId,
getToolCallUI,
getToolResultErrorMessage,
@@ -130,6 +132,14 @@ export async function handleToolEvent(
return
}
// A tool event breaks the thinking stream. Flush any open thinking
// block into contentBlocks BEFORE we add the tool_call block, or
// contentBlocks will end up with tool_call before thinking — which
// re-renders on reload in the wrong order (Mothership group above
// the Thinking block, even though thinking happened first).
flushSubagentThinkingBlock(context)
flushThinkingBlock(context)
if (isToolResultStreamEvent(event)) {
handleResultPhase(event.payload, context, parentToolCallId)
return
@@ -190,9 +200,24 @@ function handleResultPhase(
...(errorMessage ? { error: errorMessage } : {}),
endTime,
})
stampToolCallBlockEnd(context, toolCallId, endTime)
markToolResultSeen(toolCallId)
}
function stampToolCallBlockEnd(
context: StreamingContext,
toolCallId: string,
endTime: number
): void {
for (let i = context.contentBlocks.length - 1; i >= 0; i--) {
const block = context.contentBlocks[i]
if (block.type === 'tool_call' && block.toolCall?.id === toolCallId) {
if (block.endedAt === undefined) block.endedAt = endTime
return
}
}
}
async function handleCallPhase(
data: MothershipStreamV1ToolCallDescriptor,
context: StreamingContext,

View File

@@ -51,12 +51,17 @@ export function addContentBlock(
})
}
export function stampBlockEnd(block: ContentBlock): void {
if (block.endedAt === undefined) block.endedAt = Date.now()
}
/**
* Flush any open thinking block into contentBlocks and clear the thinking state.
* Safe to call repeatedly.
*/
export function flushThinkingBlock(context: StreamingContext): void {
if (context.currentThinkingBlock) {
stampBlockEnd(context.currentThinkingBlock)
context.contentBlocks.push(context.currentThinkingBlock)
}
context.isInThinkingBlock = false
@@ -65,6 +70,7 @@ export function flushThinkingBlock(context: StreamingContext): void {
export function flushSubagentThinkingBlock(context: StreamingContext): void {
if (context.currentSubagentThinkingBlock) {
stampBlockEnd(context.currentSubagentThinkingBlock)
context.contentBlocks.push(context.currentSubagentThinkingBlock)
}
context.currentSubagentThinkingBlock = null

View File

@@ -55,6 +55,7 @@ export interface ContentBlock {
toolCall?: ToolCallState
calledBy?: string
timestamp: number
endedAt?: number
}
export interface StreamingContext {