refactor: migrate thinking steps handling to new data structure and streamline related components

This commit is contained in:
Anish Sarkar 2026-03-24 02:23:05 +05:30
parent b8f3f41326
commit e587b588c9
7 changed files with 135 additions and 353 deletions

View file

@ -2,8 +2,8 @@ import type { ThreadMessageLike } from "@assistant-ui/react";
import type { MessageRecord } from "./thread-persistence";
/**
* Convert backend message to assistant-ui ThreadMessageLike format
* Filters out 'thinking-steps' part as it's handled separately via messageThinkingSteps
* Convert backend message to assistant-ui ThreadMessageLike format.
* Migrates legacy `thinking-steps` parts to `data-thinking-steps` (assistant-ui data parts).
*/
export function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike {
let content: ThreadMessageLike["content"];
@ -11,26 +11,34 @@ export function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike {
if (typeof msg.content === "string") {
content = [{ type: "text", text: msg.content }];
} else if (Array.isArray(msg.content)) {
// Filter out custom metadata parts - they're handled separately
const filteredContent = msg.content.filter((part: unknown) => {
if (typeof part !== "object" || part === null || !("type" in part)) return true;
const partType = (part as { type: string }).type;
// Filter out metadata parts not directly renderable by assistant-ui
return (
partType !== "thinking-steps" &&
partType !== "mentioned-documents" &&
partType !== "attachments"
);
});
const convertedContent = msg.content
.filter((part: unknown) => {
if (typeof part !== "object" || part === null || !("type" in part)) return true;
const partType = (part as { type: string }).type;
return partType !== "mentioned-documents" && partType !== "attachments";
})
.map((part: unknown) => {
if (
typeof part === "object" &&
part !== null &&
"type" in part &&
(part as { type: string }).type === "thinking-steps"
) {
return {
type: "data-thinking-steps",
data: { steps: (part as { steps: unknown[] }).steps ?? [] },
};
}
return part;
});
content =
filteredContent.length > 0
? (filteredContent as ThreadMessageLike["content"])
convertedContent.length > 0
? (convertedContent as ThreadMessageLike["content"])
: [{ type: "text", text: "" }];
} else {
content = [{ type: "text", text: String(msg.content) }];
}
// Build metadata.custom for author display in shared chats
const metadata = msg.author_id
? {
custom: {

View file

@ -15,6 +15,10 @@ export type ContentPart =
toolName: string;
args: Record<string, unknown>;
result?: unknown;
}
| {
type: "data-thinking-steps";
data: { steps: ThinkingStepData[] };
};
export interface ContentPartsState {
@ -23,6 +27,32 @@ export interface ContentPartsState {
toolCallIndices: Map<string, number>;
}
export function updateThinkingSteps(
state: ContentPartsState,
steps: Map<string, ThinkingStepData>
): void {
const stepsArray = Array.from(steps.values());
const existingIdx = state.contentParts.findIndex((p) => p.type === "data-thinking-steps");
if (existingIdx >= 0) {
state.contentParts[existingIdx] = {
type: "data-thinking-steps",
data: { steps: stepsArray },
};
} else {
state.contentParts.unshift({
type: "data-thinking-steps",
data: { steps: stepsArray },
});
if (state.currentTextPartIndex >= 0) {
state.currentTextPartIndex += 1;
}
for (const [id, idx] of state.toolCallIndices) {
state.toolCallIndices.set(id, idx + 1);
}
}
}
export function appendText(state: ContentPartsState, delta: string): void {
if (
state.currentTextPartIndex >= 0 &&
@ -75,6 +105,7 @@ export function buildContentForUI(
const filtered = state.contentParts.filter((part) => {
if (part.type === "text") return part.text.length > 0;
if (part.type === "tool-call") return toolsWithUI.has(part.toolName);
if (part.type === "data-thinking-steps") return true;
return false;
});
return filtered.length > 0
@ -84,23 +115,17 @@ export function buildContentForUI(
export function buildContentForPersistence(
state: ContentPartsState,
toolsWithUI: Set<string>,
currentThinkingSteps: Map<string, ThinkingStepData>
toolsWithUI: Set<string>
): unknown[] {
const parts: unknown[] = [];
if (currentThinkingSteps.size > 0) {
parts.push({
type: "thinking-steps",
steps: Array.from(currentThinkingSteps.values()),
});
}
for (const part of state.contentParts) {
if (part.type === "text" && part.text.length > 0) {
parts.push(part);
} else if (part.type === "tool-call" && toolsWithUI.has(part.toolName)) {
parts.push(part);
} else if (part.type === "data-thinking-steps") {
parts.push(part);
}
}