From 4b05e698f87910d267a93794248e2e663373de96 Mon Sep 17 00:00:00 2001 From: Shantur Rathore Date: Fri, 2 Jan 2026 16:21:24 +0000 Subject: [PATCH] Require tool part ids for tool-call rendering and caching Rebind permissions from callID to part id when parts arrive. --- packages/ui/src/components/markdown.tsx | 5 +- packages/ui/src/components/message-block.tsx | 2 +- packages/ui/src/components/tool-call.tsx | 14 ++++- packages/ui/src/stores/message-v2/bridge.ts | 43 +++++++++++-- .../src/stores/message-v2/instance-store.ts | 60 +++++++++++++++---- .../ui/src/stores/message-v2/normalizers.ts | 26 +------- 6 files changed, 107 insertions(+), 43 deletions(-) diff --git a/packages/ui/src/components/markdown.tsx b/packages/ui/src/components/markdown.tsx index 71abed0b..3379bb0b 100644 --- a/packages/ui/src/components/markdown.tsx +++ b/packages/ui/src/components/markdown.tsx @@ -35,7 +35,10 @@ export function Markdown(props: MarkdownProps) { const dark = Boolean(props.isDark) const themeKey = dark ? "dark" : "light" const highlightEnabled = !props.disableHighlight - const partId = typeof part.id === "string" && part.id.length > 0 ? part.id : "__anonymous__" + const partId = typeof part.id === "string" && part.id.length > 0 ? part.id : "" + if (!partId) { + throw new Error("Markdown rendering requires a part id") + } const versionKey = typeof part.version === "number" ? String(part.version) : "" const cacheKey = makeMarkdownCacheKey(partId, themeKey, highlightEnabled, versionKey) diff --git a/packages/ui/src/components/message-block.tsx b/packages/ui/src/components/message-block.tsx index 539909a6..fc152de6 100644 --- a/packages/ui/src/components/message-block.tsx +++ b/packages/ui/src/components/message-block.tsx @@ -453,7 +453,7 @@ export default function MessageBlock(props: MessageBlockProps) { props.toolCall) const toolName = createMemo(() => toolCallMemo()?.tool || "") - const toolCallIdentifier = createMemo(() => toolCallMemo()?.callID || props.toolCallId || toolCallMemo()?.id || "") + const toolCallIdentifier = createMemo(() => { + const partId = toolCallMemo()?.id + if (!partId) { + throw new Error("Tool call requires a part id") + } + return partId + }) const toolState = createMemo(() => toolCallMemo()?.state) const cacheContext = createMemo(() => ({ @@ -695,7 +701,11 @@ export default function ToolCall(props: ToolCallProps) { ) } - const markdownPart: TextPart = { type: "text", text: options.content, version: props.partVersion } + const partId = toolCallMemo()?.id + if (!partId) { + throw new Error("Tool call markdown requires a part id") + } + const markdownPart: TextPart = { id: partId, type: "text", text: options.content, version: props.partVersion } const cached = markdownCache.get() if (cached) { markdownPart.renderCache = cached diff --git a/packages/ui/src/stores/message-v2/bridge.ts b/packages/ui/src/stores/message-v2/bridge.ts index 5252e0f6..7c570377 100644 --- a/packages/ui/src/stores/message-v2/bridge.ts +++ b/packages/ui/src/stores/message-v2/bridge.ts @@ -112,27 +112,62 @@ function extractPermissionMessageId(permission: Permission): string | undefined } function extractPermissionPartId(permission: Permission): string | undefined { + const metadata = (permission as any).metadata || {} + return ( + (permission as any).partID || + (permission as any).partId || + metadata.partID || + metadata.partId || + undefined + ) +} + +function extractPermissionCallId(permission: Permission): string | undefined { const metadata = (permission as any).metadata || {} return ( (permission as any).callID || (permission as any).callId || (permission as any).toolCallID || (permission as any).toolCallId || - metadata.partId || - metadata.partID || metadata.callID || metadata.callId || undefined ) } +function resolvePartIdFromCallId(store: ReturnType, messageId?: string, callId?: string): string | undefined { + if (!messageId || !callId) return undefined + const record = store.getMessage(messageId) + if (!record) return undefined + for (const partId of record.partIds) { + const part = record.parts[partId]?.data + if (!part || part.type !== "tool") continue + const toolCallId = + (part as any).callID ?? + (part as any).callId ?? + (part as any).toolCallID ?? + (part as any).toolCallId ?? + undefined + if (toolCallId === callId && typeof part.id === "string" && part.id.length > 0) { + return part.id + } + } + return undefined +} + export function upsertPermissionV2(instanceId: string, permission: Permission): void { if (!permission) return const store = messageStoreBus.getOrCreate(instanceId) + const messageId = extractPermissionMessageId(permission) + let partId = extractPermissionPartId(permission) + if (!partId) { + const callId = extractPermissionCallId(permission) + partId = resolvePartIdFromCallId(store, messageId, callId) + } store.upsertPermission({ permission, - messageId: extractPermissionMessageId(permission), - partId: extractPermissionPartId(permission), + messageId, + partId, enqueuedAt: (permission as any).time?.created ?? Date.now(), }) } diff --git a/packages/ui/src/stores/message-v2/instance-store.ts b/packages/ui/src/stores/message-v2/instance-store.ts index c51eb1f6..3e3383cf 100644 --- a/packages/ui/src/stores/message-v2/instance-store.ts +++ b/packages/ui/src/stores/message-v2/instance-store.ts @@ -51,16 +51,8 @@ function ensurePartId(messageId: string, part: ClientPart, index: number): strin return part.id } - const toolCallId = - (part as any).callID ?? - (part as any).callId ?? - (part as any).toolCallID ?? - (part as any).toolCallId ?? - undefined - - if (part.type === "tool" && typeof toolCallId === "string" && toolCallId.length > 0) { - part.id = toolCallId - return toolCallId + if (part.type === "tool") { + throw new Error("Tool part missing id") } const fallbackId = `${messageId}-part-${index}` @@ -504,6 +496,50 @@ export function createInstanceMessageStore(instanceId: string, hooks?: MessageSt }) } + function rebindPermissionForPart(messageId: string, partId: string, part: ClientPart) { + if (!messageId || !partId || part.type !== "tool") { + return + } + + const toolCallId = + (part as any).callID ?? + (part as any).callId ?? + (part as any).toolCallID ?? + (part as any).toolCallId ?? + undefined + if (!toolCallId) { + return + } + + setState( + "permissions", + "byMessage", + messageId, + produce((draft) => { + if (!draft) return + const existing = draft[partId] + for (const [key, entry] of Object.entries(draft)) { + if (!entry || entry.partId) continue + const permissionCallId = + (entry.permission as any).callID ?? + (entry.permission as any).callId ?? + (entry.permission as any).toolCallID ?? + (entry.permission as any).toolCallId ?? + (entry.permission as any).metadata?.callID ?? + (entry.permission as any).metadata?.callId ?? + undefined + if (permissionCallId !== toolCallId) continue + if (!existing || existing.permission.id === entry.permission.id) { + entry.partId = partId + draft[partId] = entry + delete draft[key] + } + break + } + }), + ) + } + function applyPartUpdate(input: PartUpdateInput) { const message = state.messages[input.messageId] if (!message) { @@ -535,6 +571,8 @@ export function createInstanceMessageStore(instanceId: string, hooks?: MessageSt }), ) + rebindPermissionForPart(input.messageId, partId, cloned) + if (isCompletedTodoPart(cloned)) { recordLatestTodoSnapshot(message.sessionId, { messageId: input.messageId, @@ -740,7 +778,7 @@ export function createInstanceMessageStore(instanceId: string, hooks?: MessageSt function upsertPermission(entry: PermissionEntry) { const messageKey = entry.messageId ?? "__global__" - const partKey = entry.partId ?? "__global__" + const partKey = entry.partId ?? entry.permission?.id ?? "__global__" setState( "permissions", diff --git a/packages/ui/src/stores/message-v2/normalizers.ts b/packages/ui/src/stores/message-v2/normalizers.ts index b249fb4c..d4ddccd1 100644 --- a/packages/ui/src/stores/message-v2/normalizers.ts +++ b/packages/ui/src/stores/message-v2/normalizers.ts @@ -26,35 +26,13 @@ function decodeTextSegment(segment: any): any { return segment } -function deriveToolPartId(part: any): string | undefined { - if (!part || typeof part !== "object") { - return undefined - } - if (part.type !== "tool") { - return undefined - } - const callId = - part.callID ?? - part.callId ?? - part.toolCallID ?? - part.toolCallId ?? - undefined - if (typeof callId === "string" && callId.length > 0) { - return callId - } - return undefined -} - export function normalizeMessagePart(part: any): any { if (!part || typeof part !== "object") { return part } - if ((typeof part.id !== "string" || part.id.length === 0) && part.type === "tool") { - const inferredId = deriveToolPartId(part) - if (inferredId) { - part = { ...part, id: inferredId } - } + if (part.type === "tool" && (typeof part.id !== "string" || part.id.length === 0)) { + throw new Error("Tool part missing id") } if (part.type !== "text") {