salvacybersec/CodeNomad
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Require tool part ids for tool-call rendering and caching

Browse Source 
Rebind permissions from callID to part id when parts arrive.
This commit is contained in:
Shantur Rathore
2026-01-02 16:21:24 +00:00
parent a9524b3e30
commit 4b05e698f8
6 changed files with 107 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
packages/ui/src/components/markdown.tsx
View File

@@ -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)

2
packages/ui/src/components/message-block.tsx
View File

@@ -453,7 +453,7 @@ export default function MessageBlock(props: MessageBlockProps) {
</div>
<ToolCall
toolCall={toolItem.toolPart}
toolCallId={toolItem.key}
toolCallId={toolItem.toolPart.id}
messageId={toolItem.messageId}
messageVersion={toolItem.messageVersion}
partVersion={toolItem.partVersion}

14
packages/ui/src/components/tool-call.tsx
View File

@@ -222,7 +222,13 @@ export default function ToolCall(props: ToolCallProps) {
const { isDark } = useTheme()
const toolCallMemo = createMemo(() => 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<RenderCache>()
if (cached) {
markdownPart.renderCache = cached

43
packages/ui/src/stores/message-v2/bridge.ts
View File

@@ -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<typeof messageStoreBus.getOrCreate>, 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(),
})
}

60
packages/ui/src/stores/message-v2/instance-store.ts
View File

@@ -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",

26
packages/ui/src/stores/message-v2/normalizers.ts
View File

@@ -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") {