diff --git a/src/backend/base/langflow/schema/content_types.py b/src/backend/base/langflow/schema/content_types.py index c26282b093..dd100bd9d3 100644 --- a/src/backend/base/langflow/schema/content_types.py +++ b/src/backend/base/langflow/schema/content_types.py @@ -75,6 +75,19 @@ class BaseContent(BaseModel): ), ) + def __init__(self, **data: Any) -> None: + super().__init__(**data) + # Mark only the discriminator as "set" so partial-update callers using + # ``model_dump(exclude_unset=True)`` (notably ``aupdate_messages``) + # keep the ``type`` tag. Without it, ``TextContent(text="x")`` would + # dump to ``{"text": "x"}`` and the next read-back through + # ``MessageRead``'s discriminated union would fail with + # ``union_tag_not_found``. Other defaulted fields stay unset so true + # exclude_unset semantics survive: a patch like ``ContentBlock(title= + # "...")`` no longer overwrites an existing block's ``duration`` / + # ``header`` / ``contents`` with their defaults on merge. + self.model_fields_set.add("type") + def to_dict(self) -> dict[str, Any]: return self.model_dump() @@ -260,13 +273,10 @@ class ContentBlock(BaseContent): allow_markdown: bool = Field(default=True) media_url: list[str] | None = None - def __init__(self, **data) -> None: - super().__init__(**data) - # Mark every field as "set" so legacy callers iterating - # ``model_fields_set`` (and ``model_dump(exclude_unset=True)``) see - # the full ContentBlock shape, including the ``type="group"`` - # discriminator that downstream validators depend on. - self.model_fields_set.update(type(self).model_fields) + # ``__init__`` is inherited from ``BaseContent``: it marks only the + # ``type`` discriminator as set, so ``model_dump(exclude_unset=True)`` + # preserves ``type="group"`` while leaving ``allow_markdown`` and the + # other defaulted fields out of partial-update merges. @field_validator("contents", mode="before") @classmethod diff --git a/src/frontend/src/components/core/chatComponents/ContentBlockDisplay.tsx b/src/frontend/src/components/core/chatComponents/ContentBlockDisplay.tsx index 8fac339c1a..368432bc0e 100644 --- a/src/frontend/src/components/core/chatComponents/ContentBlockDisplay.tsx +++ b/src/frontend/src/components/core/chatComponents/ContentBlockDisplay.tsx @@ -1,14 +1,18 @@ "use client"; import { motion } from "framer-motion"; import { ChevronDown } from "lucide-react"; -import { useState } from "react"; +import { useEffect, useMemo, useRef, useState } from "react"; import { BorderTrail } from "@/components/core/border-trail"; import { useToolDurations } from "@/components/core/playgroundComponent/chat-view/chat-messages/hooks/use-tool-durations"; import { formatTime, formatToolTitle, } from "@/components/core/playgroundComponent/chat-view/chat-messages/utils/format"; -import type { ContentBlock } from "@/types/chat"; +import { + type ContentBlockItem, + type ContentType, + isGroupedBlock, +} from "@/types/chat"; import { cn } from "@/utils/utils"; import ForwardedIconComponent from "../../common/genericIconComponent"; import { @@ -19,9 +23,13 @@ import { } from "../../ui/accordion"; import ContentDisplay from "./ContentDisplay"; import DurationDisplay from "./DurationDisplay"; +import { groupConsecutiveCitations } from "./groupCitations"; +import { SourcesStrip } from "./SourcesStrip"; +import { ToolCallCard } from "./ToolCallCard"; +import { getToolStatus, TOOL_STATUS_CLASS } from "./toolStatus"; interface ContentBlockDisplayProps { - contentBlocks: ContentBlock[]; + contentBlocks: ContentBlockItem[]; isLoading?: boolean; state?: string; chatId: string; @@ -39,13 +47,62 @@ export function ContentBlockDisplay({ }: ContentBlockDisplayProps) { const [isExpanded, setIsExpanded] = useState(false); - // Use shared hook for tool duration tracking + // Separate flat content items from grouped blocks. Memoize so child + // components (notably Radix Accordion, which collects refs per render) + // see stable references across re-renders. Without this, every parent + // render produces fresh filter() arrays, every AccordionItem looks + // 'new' to the collection, refs are re-registered, and Radix loops + // until React hits the max-update-depth guard. + const groupedBlocks = useMemo( + () => contentBlocks.filter(isGroupedBlock), + [contentBlocks], + ); + const flatItems = useMemo( + () => + contentBlocks.filter( + (item): item is ContentType => !isGroupedBlock(item), + ), + [contentBlocks], + ); + + // Use shared hook for tool duration tracking (only grouped blocks) const { toolElapsedTimes, toolItems } = useToolDurations( - contentBlocks, + groupedBlocks.length > 0 ? groupedBlocks : undefined, isLoading ?? false, ); - if (!toolItems.length) { + // Stabilize the Radix Accordion defaultValue so the wrapped accordion + // doesn't see a new array reference on every parent re-render. + const runningToolKeys = useMemo( + () => + toolItems + .filter(({ content }) => getToolStatus(content) === "running") + .map(({ toolKey }) => toolKey), + [toolItems], + ); + + // Controlled open state for the tools accordion. Radix only reads + // ``defaultValue`` once at mount, so a tool that mounts or flips to + // running *after* the first render never auto-expands. Seed with the + // keys running at mount, then add any key that *newly* enters the + // running state. We only auto-open on the running transition (diffed + // against the previous set), never re-open, so a user collapse via + // ``onValueChange`` stays collapsed even while the tool is still running. + const [openToolKeys, setOpenToolKeys] = useState(runningToolKeys); + const prevRunningKeysRef = useRef(runningToolKeys); + useEffect(() => { + const newlyRunning = runningToolKeys.filter( + (key) => !prevRunningKeysRef.current.includes(key), + ); + if (newlyRunning.length > 0) { + setOpenToolKeys((prev) => + Array.from(new Set([...prev, ...newlyRunning])), + ); + } + prevRunningKeysRef.current = runningToolKeys; + }, [runningToolKeys]); + + if (!toolItems.length && !flatItems.length) { return null; } @@ -65,126 +122,184 @@ export function ContentBlockDisplay({ return (
- - {isLoading && ( - - )} - {!hideHeader && ( -
-
- {headerIcon && ( - - - - )} -

- {headerTitle} -

+ {/* Render flat content items. Three routes: + - tool_use items get wrapped in their own collapsible card + (ToolCallCard) so they always carry the bordered-header + chrome — the agent emits them flat, not inside a group + - consecutive citations get coalesced into one Sources strip + - everything else passes through ContentDisplay directly */} + {flatItems.length > 0 && ( +
+ {groupConsecutiveCitations(flatItems).map((run) => { + if (run.kind === "sources") { + return ( +
+ +
+ ); + } + if (run.item.type === "tool_use") { + return ( + + ); + } + return ( + + ); + })} +
+ )} + + {/* Render grouped blocks with accordion UI */} + {toolItems.length > 0 && ( + + {isLoading && ( + + )} + {!hideHeader && ( +
+
+ {headerIcon && ( + + + + )} +

+ {headerTitle} +

+
+
+ {!playgroundPage && ( + + )} + setIsExpanded((prev) => !prev)} + className="cursor-pointer" + > + + +
-
- {!playgroundPage && ( - - )} - setIsExpanded((prev) => !prev)} - className="cursor-pointer" + )} + + {(hideHeader || isExpanded) && ( +
+ - - -
-
- )} + {toolItems.map( + ({ content, toolKey, blockIndex, contentIndex }, flatIdx) => { + const rawTitle = + content.header?.title || + content.name || + `Tool ${flatIdx + 1}`; + const toolTitle = + typeof rawTitle === "string" + ? formatToolTitle(rawTitle) + : rawTitle; + const toolDuration = + toolElapsedTimes[toolKey] ?? content.duration ?? 0; + const status = getToolStatus(content); + const dotClass = TOOL_STATUS_CLASS[status]; + // Drop the "Called tool" / "Node" prefix and the + // bg-muted code-pill chrome — consensus across + // assistant-ui, ai-chatbot, CopilotKit, Claude, and + // ChatGPT is just the tool name in medium weight as + // the header. - {(hideHeader || isExpanded) && ( -
- - {toolItems.map( - ({ content, toolKey, blockIndex, contentIndex }, flatIdx) => { - const rawTitle = - content.header?.title || - content.name || - `Tool ${flatIdx + 1}`; - const toolTitle = - typeof rawTitle === "string" - ? formatToolTitle(rawTitle) - : rawTitle; - const isAgentStep = - content.header?.icon === "GitBranch" || - (typeof content.name === "string" && - content.name.startsWith("Node ")); - const toolLabel = isAgentStep ? "Node" : "Called tool"; - const toolDuration = - toolElapsedTimes[toolKey] ?? content.duration ?? 0; - - return ( - - -
-
-
- {toolLabel}{" "} -
-
-

+ return ( + + +

+
+ +

{toolTitle}

+ {toolDuration > 0 && ( + + {formatTime(toolDuration, true)} + + )}
-
- - {formatTime(toolDuration, true)} - + + +
+
-
- - -
- -
- - - ); - }, - )} - -
- )} - + + + ); + }, + )} + +
+ )} + + )}
); } diff --git a/src/frontend/src/components/core/chatComponents/ContentDisplay.tsx b/src/frontend/src/components/core/chatComponents/ContentDisplay.tsx index 23d3ffcc92..3c831581bd 100644 --- a/src/frontend/src/components/core/chatComponents/ContentDisplay.tsx +++ b/src/frontend/src/components/core/chatComponents/ContentDisplay.tsx @@ -1,27 +1,38 @@ -import type { ReactNode } from "react"; +import { ChevronDown } from "lucide-react"; +import { Fragment, type ReactNode, useState } from "react"; import Markdown from "react-markdown"; import rehypeMathjax from "rehype-mathjax/browser"; import remarkGfm from "remark-gfm"; -import type { ContentType, JSONValue } from "@/types/chat"; +import { formatSeconds } from "@/components/core/playgroundComponent/chat-view/chat-messages/utils/format"; +import type { ContentBlockItem, JSONValue } from "@/types/chat"; import { extractLanguage, isCodeBlock } from "@/utils/codeBlockUtils"; import ForwardedIconComponent from "../../common/genericIconComponent"; import SimplifiedCodeTabComponent from "../codeTabsComponent"; import DurationDisplay from "./DurationDisplay"; +import { SourcesStrip } from "./SourcesStrip"; +import { looksPreformatted, unwrapToolMessage } from "./toolOutput"; export default function ContentDisplay({ content, chatId, playgroundPage, }: { - content: ContentType; + // Accept any ContentBlockItem so nested ContentBlock groups land in the + // `case "group"` branch below instead of falling through to no rendering. + content: ContentBlockItem; chatId: string; playgroundPage?: boolean; }) { - const renderDuration = content.duration !== undefined && !playgroundPage && ( -
- -
- ); + // Reasoning blocks surface their own duration inline via ReasoningDisplay's + // "Thought for Xs" label, so skip the absolute top-right DurationDisplay + // there to avoid rendering the same duration twice. + const renderDuration = content.duration !== undefined && + content.type !== "reasoning" && + !playgroundPage && ( +
+ +
+ ); // Then render the specific content based on type let contentData: ReactNode | null = null; @@ -127,11 +138,23 @@ export default function ContentDisplay({ break; case "tool_use": { - const formatToolOutput = (output: JSONValue) => { - if (output === null || output === undefined) return ""; + // Tool output rendering routes by the *unwrapped* payload shape: + // - markdown-y string -> Markdown renderer (prose, no chrome) + // - pre-formatted text -> code tab with language=text + // (monospace, contained horizontal scroll, + // built-in copy button) + // - object / array -> code tab with language=json (same) + // The LangChain ToolMessage envelope is unwrapped first so the + // readable `content` field becomes the body and the plumbing + // metadata (already shown by the accordion trigger) stays hidden. + const formatToolOutput = (raw: JSONValue) => { + const output = unwrapToolMessage(raw); + if (output === null || output === undefined) return null; - // If it's a string, render as markdown if (typeof output === "string") { + if (looksPreformatted(output)) { + return ; + } return ( 0; + // Treat an empty (or whitespace-only) string as "no output" so it + // doesn't render an empty bordered box. Legitimate falsy outputs like + // 0 or false still render through formatToolOutput. + const hasOutput = + content.output !== undefined && + content.output !== null && + !(typeof content.output === "string" && content.output.trim() === ""); + const hasError = content.error != null; + // Eyebrow labels (INPUT/OUTPUT/ERROR) used to bracket each section, + // but the surrounding accordion card is already the "tool call" + // context — extra labels just add chrome. Match the assistant-ui / + // Claude pattern: args and result stack directly inside the card, + // separated by a hairline rule. Empty sections render nothing. + const showSeparator = hasInput && (hasOutput || hasError); contentData = ( -
- - **Input:** - - - {content.output !== undefined && ( - <> - - **Output:** - -
{formatToolOutput(content.output)}
- +
+ {hasInput && } + {showSeparator &&
} + {hasOutput && ( +
+ {formatToolOutput(content.output as JSONValue)} +
)} - {content.error != null && ( -
- - **Error:** - + {hasError && ( +
); break; + + case "image": + contentData = ( +
+ {content.urls?.map((url, index) => ( + {content.caption + ))} + {/* base64 is a fallback for when no usable URL is provided. + `some(Boolean)` so urls=[""] or urls=[null] don't suppress the + fallback while rendering a broken above. */} + {!content.urls?.some(Boolean) && content.base64 && ( + {content.caption + )} + {content.caption && ( +

{content.caption}

+ )} +
+ ); + break; + + case "audio": + contentData = ( +
+ {content.urls?.map((url, index) => ( + + ))} + {!content.urls?.some(Boolean) && content.base64 && ( + + )} + {content.transcript && ( +

+ {content.transcript} +

+ )} +
+ ); + break; + + case "video": + contentData = ( +
+ {content.urls?.map((url, index) => ( + + ))} + {!content.urls?.some(Boolean) && content.base64 && ( + + )} +
+ ); + break; + + case "file": + contentData = ( +
+ + {content.urls?.map((url, index) => ( + + {content.filename || `Download file ${index + 1}`} + + ))} +
+ ); + break; + + case "reasoning": + contentData = ( + + ); + break; + + case "usage": { + // Backend serializes Optional[int] as JSON null, so `!= null` catches + // both undefined and null. Using `!== undefined` here would render the + // literal string "null in / null out" when the producer reports no + // counts. + const hasInput = content.input_tokens != null; + const hasOutput = content.output_tokens != null; + const hasTokens = hasInput || hasOutput; + contentData = ( +
+ {content.model && ( + {content.model} + )} + {hasTokens && ( + + Tokens: {hasInput ? `${content.input_tokens} in` : ""} + {hasInput && hasOutput ? " / " : ""} + {hasOutput ? `${content.output_tokens} out` : ""} + + )} +
+ ); + break; + } + + case "citation": + // Single citation renders as a one-card Sources strip; consecutive + // flat citations are coalesced into a multi-card strip by + // ContentBlockDisplay before they reach this branch. + contentData = ; + break; + + case "group": + // A nested ContentBlock inside another container. Render the title as + // a small section header and recurse on each child. The outer + // ContentBlockDisplay handles top-level groups; this branch only + // fires when a producer nests a group inside a parent's contents. + contentData = ( +
+ {content.title && ( +

{content.title}

+ )} + {content.contents?.map((child, index) => ( + + ))} +
+ ); + break; } return ( @@ -251,3 +426,87 @@ export default function ContentDisplay({
); } + +/** Renders a tool's input. Flat objects (string/number/bool/null values + * only) become labelled rows; anything nested falls back to a JSON block + * so the structure stays readable. Empty input is handled by the caller + * (the surrounding card simply skips this component). */ +function ToolInputDisplay({ input }: { input: Record }) { + const entries = Object.entries(input); + // A value counts as "flat" if it's a primitive or a list of primitives. + // Nested objects (or arrays of objects) fall through to the JSON block + // where indentation makes them readable. + const isPrimitive = (v: JSONValue) => v === null || typeof v !== "object"; + const isFlat = entries.every( + ([, v]) => isPrimitive(v) || (Array.isArray(v) && v.every(isPrimitive)), + ); + if (!isFlat) { + return ( + + ); + } + return ( +
+ {entries.map(([key, value]) => ( + +
{key}
+
{JSON.stringify(value)}
+ + ))} +
+ ); +} + +/** Reasoning section. Live shimmer while streaming, collapsible summary once + * the producer attaches a `duration` (its signal that the step is done). */ +function ReasoningDisplay({ + text, + duration, +}: { + text: string; + duration?: number; +}) { + const [isOpen, setIsOpen] = useState(false); + const isStreaming = duration === undefined; + + if (isStreaming) { + return ( +
+ + Thinking… +
+ ); + } + + return ( +
+ + {isOpen && ( +
+ {text} +
+ )} +
+ ); +} diff --git a/src/frontend/src/components/core/chatComponents/SourcesStrip.tsx b/src/frontend/src/components/core/chatComponents/SourcesStrip.tsx new file mode 100644 index 0000000000..6496dbad84 --- /dev/null +++ b/src/frontend/src/components/core/chatComponents/SourcesStrip.tsx @@ -0,0 +1,100 @@ +import type { CitationContent } from "@/types/chat"; + +/** Returns the URL if its scheme is http(s), otherwise null. Guards against + * `javascript:` / `data:` / `vbscript:` and other unsafe schemes that + * React would otherwise pass through to a clickable anchor in production. */ +function safeUrl(raw: string | null | undefined): string | null { + if (!raw) return null; + try { + const parsed = new URL(raw); + return parsed.protocol === "http:" || parsed.protocol === "https:" + ? raw + : null; + } catch { + return null; + } +} + +/** Best-effort host extraction. Returns null for unparseable inputs so the + * caller can decide whether to fall back to the raw URL or hide the chip. */ +function extractDomain(raw: string | null | undefined): string | null { + if (!raw) return null; + try { + const parsed = new URL(raw); + return parsed.hostname.replace(/^www\./, ""); + } catch { + return null; + } +} + +function FaviconImg({ domain }: { domain: string }) { + // Google's favicon service handles missing icons gracefully (returns a + // globe), so we don't need to plumb our own onError state. + const src = `https://www.google.com/s2/favicons?domain=${encodeURIComponent(domain)}&sz=32`; + return ( + + ); +} + +function SourceCard({ citation }: { citation: CitationContent }) { + const url = safeUrl(citation.url); + const domain = extractDomain(citation.url) ?? "source"; + const heading = citation.title || domain; + const inner = ( +
+
+ + {domain} +
+

{heading}

+ {citation.cited_text && ( +

+ {citation.cited_text} +

+ )} +
+ ); + if (!url) { + // URL was missing or unsafe — render the card content as a div so + // there's no clickable surface routing to a javascript: handler. + return inner; + } + return ( + + {inner} + + ); +} + +/** Horizontally-scrolling row of source cards. Used both for single + * citations and for coalesced groups of consecutive citations. */ +export function SourcesStrip({ citations }: { citations: CitationContent[] }) { + if (citations.length === 0) return null; + return ( +
+
+ {citations.length > 1 ? "Sources" : "Source"} +
+
+ {citations.map((citation, idx) => ( + + ))} +
+
+ ); +} diff --git a/src/frontend/src/components/core/chatComponents/ToolCallCard.tsx b/src/frontend/src/components/core/chatComponents/ToolCallCard.tsx new file mode 100644 index 0000000000..c62d2c1023 --- /dev/null +++ b/src/frontend/src/components/core/chatComponents/ToolCallCard.tsx @@ -0,0 +1,102 @@ +import { useEffect, useState } from "react"; +import { + formatTime, + formatToolTitle, +} from "@/components/core/playgroundComponent/chat-view/chat-messages/utils/format"; +import type { ToolContent } from "@/types/chat"; +import { cn } from "@/utils/utils"; +import { + Accordion, + AccordionContent, + AccordionItem, + AccordionTrigger, +} from "../../ui/accordion"; +import ContentDisplay from "./ContentDisplay"; +import { getToolStatus, TOOL_STATUS_CLASS } from "./toolStatus"; + +/** Bordered collapsible card for a single tool call. Matches the + * assistant-ui / ai-chatbot / CopilotKit / Claude / ChatGPT consensus: + * header is a clickable row with a status dot, the tool name in + * medium-weight mono, and a duration on the right. Body holds the args + * and result via ContentDisplay's tool_use branch. + * + * Auto-expands while the tool is running; collapses to header-only once + * the producer attaches a duration (the resolved state). User can toggle + * manually either way after that. */ +export function ToolCallCard({ + content, + chatId, + playgroundPage, +}: { + content: ToolContent; + chatId: string; + playgroundPage?: boolean; +}) { + const status = getToolStatus(content); + const dotClass = TOOL_STATUS_CLASS[status]; + + const rawTitle = content.header?.title || content.name || "Tool"; + const toolTitle = + typeof rawTitle === "string" ? formatToolTitle(rawTitle) : rawTitle; + const toolDuration = content.duration ?? 0; + const itemKey = chatId; + + // Controlled open state. Radix reads ``defaultValue`` only once at mount, + // so a card that mounts collapsed and later flips to running (history + // replay, scrollback) never auto-expands. Drive ``value`` from state and + // open on the running transition; the deps guard means a user collapse + // while the tool is still running isn't re-opened on the next render. + const [value, setValue] = useState(status === "running" ? itemKey : ""); + useEffect(() => { + if (status === "running") { + setValue(itemKey); + } + }, [status, itemKey]); + + return ( + + + +
+
+ +

+ {toolTitle} +

+
+ {toolDuration > 0 && ( + + {formatTime(toolDuration, true)} + + )} +
+ + +
+ +
+ + + + ); +} diff --git a/src/frontend/src/components/core/chatComponents/__tests__/ContentDisplay.test.tsx b/src/frontend/src/components/core/chatComponents/__tests__/ContentDisplay.test.tsx new file mode 100644 index 0000000000..7aa1df407b --- /dev/null +++ b/src/frontend/src/components/core/chatComponents/__tests__/ContentDisplay.test.tsx @@ -0,0 +1,399 @@ +import { render, screen } from "@testing-library/react"; +import type { ContentBlockItem } from "@/types/chat"; +import ContentDisplay from "../ContentDisplay"; + +// react-markdown ships ESM that jest doesn't transpile by default; the +// MarkdownComponent only matters for text/error/json/code cases, none of +// which these tests touch. +jest.mock("react-markdown", () => ({ + __esModule: true, + default: ({ children }: { children: React.ReactNode }) => <>{children}, +})); +jest.mock("rehype-mathjax/browser", () => () => {}); +jest.mock("remark-gfm", () => () => {}); +// Mocks use module identifiers as the importer (ContentDisplay) sees them, +// which is why these paths are relative to ContentDisplay.tsx (one dir up +// from this test file), not to the test file itself. +jest.mock("@/components/common/genericIconComponent", () => ({ + __esModule: true, + default: ({ name }: { name: string }) => ( + + ), + ForwardedIconComponent: ({ name }: { name: string }) => ( + + ), +})); +jest.mock("@/components/core/codeTabsComponent", () => ({ + __esModule: true, + default: () =>
, +})); +jest.mock("../DurationDisplay", () => ({ + __esModule: true, + default: () =>
, +})); + +describe("ContentDisplay", () => { + describe("usage", () => { + it("hides the 'Tokens:' label when both counts are null", () => { + // Backend Optional[int] serializes as null when absent. Regression + // guard: previously rendered the literal string 'null in / null out' + // because the check used `!== undefined`. + const usage = { + type: "usage", + model: "gpt-4o", + input_tokens: null, + output_tokens: null, + } as unknown as ContentBlockItem; + render(); + expect(screen.queryByText(/Tokens:/)).not.toBeInTheDocument(); + expect(screen.queryByText(/null/)).not.toBeInTheDocument(); + expect(screen.getByText("gpt-4o")).toBeInTheDocument(); + }); + + it("renders both counts when they are zero", () => { + // 0 is a legitimate count; `!= null` keeps it, `!== undefined` would + // also keep it. This is a forward-compat guard for future refactors. + const usage = { + type: "usage", + model: "claude", + input_tokens: 0, + output_tokens: 0, + } as unknown as ContentBlockItem; + render(); + expect(screen.getByText(/0 in/)).toBeInTheDocument(); + expect(screen.getByText(/0 out/)).toBeInTheDocument(); + }); + + it("renders only the input count when output is null", () => { + const usage = { + type: "usage", + input_tokens: 42, + output_tokens: null, + } as unknown as ContentBlockItem; + render(); + const span = screen.getByText(/Tokens:/); + expect(span).toHaveTextContent("Tokens: 42 in"); + expect(span).not.toHaveTextContent("null"); + }); + }); + + describe("image", () => { + it("renders the base64 fallback when urls contains only empty strings", () => { + // Regression guard: previously suppressed base64 whenever urls had any + // length, so urls=[''] killed the working fallback. + const image = { + type: "image", + urls: [""], + base64: "AAAA", + mime_type: "image/png", + caption: "fallback", + } as unknown as ContentBlockItem; + render(); + const imgs = screen.getAllByRole("img"); + const base64Img = imgs.find((img) => + (img as HTMLImageElement).src.startsWith("data:image/png;base64,"), + ); + expect(base64Img).toBeDefined(); + }); + + it("does not duplicate the image when a valid URL and base64 are both set", () => { + const image = { + type: "image", + urls: ["https://example.com/a.png"], + base64: "AAAA", + mime_type: "image/png", + } as unknown as ContentBlockItem; + render(); + const imgs = screen.getAllByRole("img"); + const base64Img = imgs.find((img) => + (img as HTMLImageElement).src.startsWith("data:"), + ); + expect(base64Img).toBeUndefined(); + expect(imgs).toHaveLength(1); + }); + }); + + describe("reasoning", () => { + it("shows a 'Thinking…' shimmer while the duration is absent", () => { + // No duration on the content means the producer is still emitting + // reasoning chunks. The renderer should show a live indicator + // rather than the resolved summary. + const reasoning = { + type: "reasoning", + text: "Considering options...", + } as unknown as ContentBlockItem; + render(); + const label = screen.getByText(/Thinking/i); + expect(label).toBeInTheDocument(); + // Live label gets animate-pulse so the user sees it as in-flight. + expect(label.className).toMatch(/animate-pulse/); + }); + + it("collapses by default with a 'Thought for …' summary when duration is set", () => { + const reasoning = { + type: "reasoning", + text: "I checked the docs and decided to call the weather tool.", + duration: 3200, + } as unknown as ContentBlockItem; + render(); + expect(screen.getByText(/Thought for/)).toBeInTheDocument(); + // Body stays collapsed until the user clicks the summary trigger. + expect(screen.queryByText(/I checked the docs/)).not.toBeInTheDocument(); + }); + }); + + describe("renderDuration", () => { + it("renders the top-right duration for a non-reasoning block", () => { + const tool = { + type: "tool_use", + name: "search", + tool_input: {}, + duration: 1200, + } as unknown as ContentBlockItem; + render(); + expect(screen.getByTestId("duration")).toBeInTheDocument(); + }); + + it("skips the top-right duration for reasoning to avoid double render", () => { + // ReasoningDisplay already surfaces the duration inline as + // "Thought for Xs"; the absolute top-right DurationDisplay must be + // skipped for reasoning so it doesn't render twice. + const reasoning = { + type: "reasoning", + text: "Decided to call the weather tool.", + duration: 3200, + } as unknown as ContentBlockItem; + render(); + expect(screen.getByText(/Thought for/)).toBeInTheDocument(); + expect(screen.queryByTestId("duration")).not.toBeInTheDocument(); + }); + }); + + describe("tool_use", () => { + it("reads input from the 'input' alias when the backend dumped by_alias=True", () => { + // The Python ToolContent.tool_input field has alias="input"; AG-UI + // and other serialization paths emit by_alias=True, so the stored + // shape uses `input` not `tool_input`. Renderer must fall back. + const tool = { + type: "tool_use", + name: "fetch_content", + input: { urls: ["https://langflow.org"] }, + } as unknown as ContentBlockItem; + render(); + expect(screen.getByText("urls")).toBeInTheDocument(); + expect(screen.queryByText(/no arguments/)).not.toBeInTheDocument(); + }); + + it("renders flat input as key/value rows, not JSON", () => { + // {query: "weather", units: "metric"} is a flat object — easier to + // scan as two rows than as JSON with braces and quotes. + const tool = { + type: "tool_use", + name: "weather", + tool_input: { query: "weather", units: "metric" }, + } as unknown as ContentBlockItem; + render(); + expect(screen.getByText("query")).toBeInTheDocument(); + expect(screen.getByText('"weather"')).toBeInTheDocument(); + expect(screen.getByText("units")).toBeInTheDocument(); + expect(screen.getByText('"metric"')).toBeInTheDocument(); + // No JSON code block for the flat case. + expect(screen.queryByTestId("code-tabs")).not.toBeInTheDocument(); + }); + + it("falls back to JSON block when input has nested values", () => { + // Nested object would be ugly in row form — code block keeps it + // readable and copyable. + const tool = { + type: "tool_use", + name: "search", + tool_input: { filters: { min: 1, max: 10 } }, + } as unknown as ContentBlockItem; + render(); + expect(screen.getByTestId("code-tabs")).toBeInTheDocument(); + }); + + it("unwraps a LangChain ToolMessage-shaped output to render only its content", () => { + // Backend tools return {content, name, id, tool_call_id, status} — + // the LangChain ToolMessage shape. The metadata is duplicated by + // the accordion trigger; only `content` is interesting to show. + const tool = { + type: "tool_use", + name: "fetch_content", + tool_input: {}, + output: { + content: "the readable body", + name: "fetch_content", + id: "abc-123", + tool_call_id: "toolu_xyz", + status: "success", + }, + } as unknown as ContentBlockItem; + render(); + expect(screen.getByText(/the readable body/)).toBeInTheDocument(); + // Plumbing metadata should not be rendered. + expect(screen.queryByText(/toolu_xyz/)).not.toBeInTheDocument(); + expect(screen.queryByText(/abc-123/)).not.toBeInTheDocument(); + }); + + it("routes pre-formatted string output through the code tab (not markdown)", () => { + // Pandas df.to_string() style output has wide column padding that + // markdown rendering would mangle. Detection should kick it into + // the monospace code-tab path instead. + const tool = { + type: "tool_use", + name: "df_describe", + tool_input: {}, + output: + " text url\n" + + "0 Langflow | Low-code AI builder for agentic... https://langflow.org", + } as unknown as ContentBlockItem; + render(); + expect(screen.getByTestId("code-tabs")).toBeInTheDocument(); + }); + + it("hides the OUTPUT section entirely when output is null", () => { + // output: null still rendered a visible 'OUTPUT' eyebrow with no + // body. Treat null the same as undefined so the empty section + // doesn't render at all. + const tool = { + type: "tool_use", + name: "search", + tool_input: {}, + output: null, + } as unknown as ContentBlockItem; + render(); + expect(screen.queryByText("OUTPUT")).not.toBeInTheDocument(); + }); + + it("hides the output section when output is an empty string", () => { + // An empty (or whitespace-only) string output used to render an empty + // bordered box. Treat it as no output so nothing renders. The output + // box is the only `.max-h-96 overflow-auto` element here. + const tool = { + type: "tool_use", + name: "search", + tool_input: {}, + output: " ", + } as unknown as ContentBlockItem; + const { container } = render( + , + ); + expect(container.querySelector(".max-h-96")).toBeNull(); + }); + + it("keeps the JSON block when output has extra fields beyond ToolMessage", () => { + // Don't unwrap if the producer added fields the user might need — + // could be tool-specific data the renderer shouldn't drop silently. + const tool = { + type: "tool_use", + name: "fetch_content", + tool_input: {}, + output: { + content: "the body", + custom_field: 42, + }, + } as unknown as ContentBlockItem; + render(); + // Falls through to the JSON code block (mocked as code-tabs). + expect(screen.getByTestId("code-tabs")).toBeInTheDocument(); + }); + + it("renders the error body in a destructive-toned panel", () => { + // No eyebrow label — the surrounding accordion already says + // "tool call" and the destructive bg + color carries the meaning. + const tool = { + type: "tool_use", + name: "search", + tool_input: {}, + error: "boom", + } as unknown as ContentBlockItem; + const { container } = render( + , + ); + // The destructive panel wraps the code tab. + expect(container.querySelector(".bg-destructive\\/10")).toBeTruthy(); + }); + + it("renders nothing for empty input (no 'no arguments' placeholder)", () => { + // The surrounding accordion provides enough context; an explicit + // empty-state line just adds chrome. The OUTPUT below carries the + // useful information. + const tool = { + type: "tool_use", + name: "search", + tool_input: {}, + output: "hi", + } as unknown as ContentBlockItem; + render(); + expect(screen.queryByText(/no arguments/)).not.toBeInTheDocument(); + }); + }); + + describe("citation", () => { + it("renders the domain extracted from the URL alongside the title", () => { + const citation = { + type: "citation", + url: "https://docs.python.org/3/library/typing.html", + title: "typing — Support for type hints", + } as unknown as ContentBlockItem; + render(); + expect(screen.getByText(/docs\.python\.org/)).toBeInTheDocument(); + expect( + screen.getByText("typing — Support for type hints"), + ).toBeInTheDocument(); + }); + + it("does not render an anchor for non-http URLs (sanitization)", () => { + // javascript: and other non-http(s) schemes must never become a + // clickable link; we degrade to a text-only card instead. + const citation = { + type: "citation", + url: "javascript:alert(1)", + title: "bad", + } as unknown as ContentBlockItem; + render(); + expect(screen.queryByRole("link")).not.toBeInTheDocument(); + expect(screen.getByText("bad")).toBeInTheDocument(); + }); + + it("renders the cited_text snippet as plain text, not HTML", () => { + // The cited_text comes from upstream model output and may contain + // angle brackets. React text nodes escape them, so the literal + // characters render rather than parsing as markup. + const citation = { + type: "citation", + url: "https://example.com", + title: "src", + cited_text: "", + } as unknown as ContentBlockItem; + render(); + expect(screen.getByText("")).toBeInTheDocument(); + }); + }); + + describe("group", () => { + it("renders the title and recurses through nested contents", () => { + // Nested ContentBlock inside another container: previously fell + // through ContentDisplay's switch and rendered nothing. Two visible + // children prove recursion: usage and citation both render text + // directly without a reveal-on-click behavior. + const nested = { + type: "group", + title: "Inner step", + contents: [ + { type: "usage", input_tokens: 1, output_tokens: 2 }, + { + type: "citation", + url: "https://example.com", + title: "Cited source", + }, + ], + } as unknown as ContentBlockItem; + render(); + expect(screen.getByText("Inner step")).toBeInTheDocument(); + expect(screen.getByText(/Tokens: 1 in \/ 2 out/)).toBeInTheDocument(); + expect(screen.getByText("Cited source")).toBeInTheDocument(); + }); + }); +}); diff --git a/src/frontend/src/components/core/chatComponents/__tests__/groupCitations.test.ts b/src/frontend/src/components/core/chatComponents/__tests__/groupCitations.test.ts new file mode 100644 index 0000000000..ddb336c4b5 --- /dev/null +++ b/src/frontend/src/components/core/chatComponents/__tests__/groupCitations.test.ts @@ -0,0 +1,49 @@ +import type { ContentType } from "@/types/chat"; +import { groupConsecutiveCitations } from "../groupCitations"; + +const text = (s: string) => + ({ type: "text", text: s }) as unknown as ContentType; +const cite = (url: string) => + ({ type: "citation", url }) as unknown as ContentType; + +describe("groupConsecutiveCitations", () => { + it("emits each non-citation item as its own single run", () => { + const runs = groupConsecutiveCitations([text("a"), text("b")]); + expect(runs).toHaveLength(2); + expect(runs.every((r) => r.kind === "single")).toBe(true); + }); + + it("coalesces a run of consecutive citations into one sources run", () => { + const runs = groupConsecutiveCitations([ + cite("https://a"), + cite("https://b"), + cite("https://c"), + ]); + expect(runs).toHaveLength(1); + expect(runs[0].kind).toBe("sources"); + if (runs[0].kind === "sources") { + expect(runs[0].citations).toHaveLength(3); + } + }); + + it("keeps non-adjacent citation groups separate", () => { + // text in the middle splits the run into two Sources groups. + const runs = groupConsecutiveCitations([ + cite("https://a"), + text("answer"), + cite("https://b"), + ]); + expect(runs).toHaveLength(3); + expect(runs[0].kind).toBe("sources"); + expect(runs[1].kind).toBe("single"); + expect(runs[2].kind).toBe("sources"); + }); + + it("still wraps a single citation in a sources run for consistent UI", () => { + // Single-citation case must still flow through the strip so the + // visual treatment is uniform. + const runs = groupConsecutiveCitations([cite("https://a")]); + expect(runs).toHaveLength(1); + expect(runs[0].kind).toBe("sources"); + }); +}); diff --git a/src/frontend/src/components/core/chatComponents/__tests__/toolOutput.test.ts b/src/frontend/src/components/core/chatComponents/__tests__/toolOutput.test.ts new file mode 100644 index 0000000000..133298984d --- /dev/null +++ b/src/frontend/src/components/core/chatComponents/__tests__/toolOutput.test.ts @@ -0,0 +1,86 @@ +import type { JSONValue } from "@/types/chat"; +import { looksPreformatted, unwrapToolMessage } from "../toolOutput"; + +describe("unwrapToolMessage", () => { + it("unwraps a canonical LangChain ToolMessage shape down to its .content", () => { + const out = unwrapToolMessage({ + content: "the body", + name: "fetch", + id: "abc", + tool_call_id: "toolu_x", + status: "success", + }); + expect(out).toBe("the body"); + }); + + it("leaves the envelope intact when extra fields are present", () => { + // The producer might be carrying tool-specific data; silently + // dropping it would lose information. + const envelope = { content: "the body", custom_field: 42 } as JSONValue; + expect(unwrapToolMessage(envelope)).toEqual(envelope); + }); + + it("leaves arrays untouched", () => { + const arr = [1, 2, 3] as JSONValue; + expect(unwrapToolMessage(arr)).toBe(arr); + }); + + it("leaves objects without a content key untouched", () => { + const obj = { result: 42, count: 7 } as JSONValue; + expect(unwrapToolMessage(obj)).toBe(obj); + }); + + it("leaves primitives untouched", () => { + expect(unwrapToolMessage("hello")).toBe("hello"); + expect(unwrapToolMessage(42)).toBe(42); + expect(unwrapToolMessage(null)).toBe(null); + }); +}); + +describe("looksPreformatted", () => { + it("flags pandas-style column-padded output", () => { + // The actual pandas df.to_string() output that broke the playground: + // dozens of consecutive spaces between column values. + const s = + " text url\n" + + "0 Langflow | Low-code AI builder for agentic and... https://langflow.org"; + expect(looksPreformatted(s)).toBe(true); + }); + + it("flags strings containing tabs", () => { + expect(looksPreformatted("col1\tcol2\tcol3")).toBe(true); + }); + + it("flags strings with very long single lines", () => { + // A 130-char single line is almost certainly a log entry or wide + // payload that markdown would mangle. + expect(looksPreformatted("x".repeat(130))).toBe(true); + }); + + it("leaves ordinary prose alone", () => { + expect( + looksPreformatted( + "Langflow.org is up — it's described as a Low-code AI builder.", + ), + ).toBe(false); + }); + + it("leaves short markdown alone", () => { + expect(looksPreformatted("**Result:** the answer is 42")).toBe(false); + }); + + it("ignores leading-indent CommonMark code blocks", () => { + // 4-space-indented lines are a CommonMark indented code block. The old + // `/ {4,}/` heuristic tripped on the leading run and forced the whole + // output into a monospace block; markdown should still render here. + const s = + "Here is the code:\n\n const x = 1;\n const y = 2;\n\nDone."; + expect(looksPreformatted(s)).toBe(false); + }); + + it("still flags interior column padding (table alignment)", () => { + // 4+ spaces AFTER content on a line is genuine column alignment that + // markdown would mangle. + expect(looksPreformatted("name value\nfoo 42")).toBe(true); + }); +}); diff --git a/src/frontend/src/components/core/chatComponents/__tests__/toolStatus.test.ts b/src/frontend/src/components/core/chatComponents/__tests__/toolStatus.test.ts new file mode 100644 index 0000000000..849101a3ce --- /dev/null +++ b/src/frontend/src/components/core/chatComponents/__tests__/toolStatus.test.ts @@ -0,0 +1,47 @@ +import type { ContentType } from "@/types/chat"; +import { getToolStatus } from "../toolStatus"; + +describe("getToolStatus", () => { + it("returns 'running' when a tool has no duration yet", () => { + const tool = { + type: "tool_use", + name: "search", + tool_input: {}, + } as unknown as ContentType; + expect(getToolStatus(tool)).toBe("running"); + }); + + it("returns 'done' once a tool reports a duration", () => { + const tool = { + type: "tool_use", + name: "search", + tool_input: {}, + duration: 1200, + } as unknown as ContentType; + expect(getToolStatus(tool)).toBe("done"); + }); + + it("returns 'error' when the tool reports an error, even with a duration", () => { + // Error wins: an errored tool shouldn't look successful just because + // the producer also stamped a duration on the failed call. + const tool = { + type: "tool_use", + name: "search", + tool_input: {}, + duration: 800, + error: "boom", + } as unknown as ContentType; + expect(getToolStatus(tool)).toBe("error"); + }); + + it("returns 'done' for non-tool_use content with a duration", () => { + // Other content types (reasoning, etc.) carrying a duration count + // as resolved even though they have no notion of 'error'. + const reasoning = { + type: "reasoning", + text: "thought", + duration: 500, + } as unknown as ContentType; + expect(getToolStatus(reasoning)).toBe("done"); + }); +}); diff --git a/src/frontend/src/components/core/chatComponents/groupCitations.ts b/src/frontend/src/components/core/chatComponents/groupCitations.ts new file mode 100644 index 0000000000..22f3cc856c --- /dev/null +++ b/src/frontend/src/components/core/chatComponents/groupCitations.ts @@ -0,0 +1,35 @@ +import type { CitationContent, ContentType } from "@/types/chat"; + +/** A run of flat content items: either a single non-citation item rendered + * via ContentDisplay, or a coalesced run of citations rendered as one + * SourcesStrip. The synthetic 'sources' kind exists so the renderer can + * group consecutive citations into one Sources row without inventing a + * fake content type or mutating the original list. */ +export type FlatRun = + | { kind: "single"; item: ContentType; index: number } + | { kind: "sources"; citations: CitationContent[]; index: number }; + +/** Walk the flat items list and merge consecutive citation entries into + * one 'sources' run. Non-citation items pass through as 'single' runs. + * `index` is the position of the first contributing item, used as a + * stable React key. */ +export function groupConsecutiveCitations(items: ContentType[]): FlatRun[] { + const runs: FlatRun[] = []; + let i = 0; + while (i < items.length) { + const item = items[i]; + if (item.type !== "citation") { + runs.push({ kind: "single", item, index: i }); + i += 1; + continue; + } + const start = i; + const citations: CitationContent[] = []; + while (i < items.length && items[i].type === "citation") { + citations.push(items[i] as CitationContent); + i += 1; + } + runs.push({ kind: "sources", citations, index: start }); + } + return runs; +} diff --git a/src/frontend/src/components/core/chatComponents/toolOutput.ts b/src/frontend/src/components/core/chatComponents/toolOutput.ts new file mode 100644 index 0000000000..cd4d327ab3 --- /dev/null +++ b/src/frontend/src/components/core/chatComponents/toolOutput.ts @@ -0,0 +1,51 @@ +import type { JSONValue } from "@/types/chat"; + +/** LangChain ToolMessage envelope shape — the readable payload lives on + * `.content`; the rest is plumbing already exposed by the surrounding + * accordion trigger (tool name, id, status). Used to decide whether to + * unwrap an output object down to its `.content` field. */ +const TOOL_MESSAGE_KEYS = new Set([ + "content", + "name", + "id", + "tool_call_id", + "status", +]); + +/** Strip the LangChain ToolMessage envelope from a tool output, returning + * the inner `content` value when the keys are a subset of the canonical + * set. Any extra key (custom tool-specific data) prevents unwrap so the + * user doesn't silently lose information. */ +export function unwrapToolMessage(output: JSONValue): JSONValue { + if ( + output !== null && + typeof output === "object" && + !Array.isArray(output) && + "content" in output && + Object.keys(output).every((k) => TOOL_MESSAGE_KEYS.has(k)) + ) { + return (output as Record).content; + } + return output; +} + +/** Detect strings that look like pre-formatted output (pandas + * df.to_string(), ASCII tables, fixed-width logs) rather than prose or + * markdown. Markdown rendering mangles column alignment, so these need to + * render in a monospace 
 block with horizontal scroll instead.
+ *
+ * Heuristics, any one fires:
+ *  - 4+ spaces after non-space content on a line (interior column padding)
+ *  - a line longer than 120 chars (long log line or wide table row)
+ *  - tab characters (very rare in markdown, common in tool output)
+ */
+export function looksPreformatted(s: string): boolean {
+  if (s.includes("\t")) return true;
+  // Interior column padding (4+ spaces after non-space content on a line)
+  // signals a fixed-width table or aligned output. A *leading* run of 4+
+  // spaces is just a CommonMark indented code block, which should still go
+  // through the markdown renderer, so require a preceding non-space char.
+  if (/\S {4,}/.test(s)) return true;
+  const longestLine = s.split("\n").reduce((m, l) => Math.max(m, l.length), 0);
+  return longestLine > 120;
+}
diff --git a/src/frontend/src/components/core/chatComponents/toolStatus.ts b/src/frontend/src/components/core/chatComponents/toolStatus.ts
new file mode 100644
index 0000000000..9a9004b5ae
--- /dev/null
+++ b/src/frontend/src/components/core/chatComponents/toolStatus.ts
@@ -0,0 +1,24 @@
+import type { ContentType } from "@/types/chat";
+
+export type ToolStatus = "running" | "done" | "error";
+
+/** Derives a tool's visible status from the data the producer already sets.
+ * An `error` payload wins (even with a duration, an errored tool shouldn't
+ * look successful); a `duration` means the step resolved; otherwise the
+ * call is still in flight. */
+export function getToolStatus(content: ContentType): ToolStatus {
+  if (content.type === "tool_use" && content.error != null) {
+    return "error";
+  }
+  if (content.duration !== undefined) {
+    return "done";
+  }
+  return "running";
+}
+
+/** Tailwind classes for the small status dot on the accordion trigger. */
+export const TOOL_STATUS_CLASS: Record = {
+  running: "bg-primary animate-pulse",
+  done: "bg-accent-emerald-foreground",
+  error: "bg-destructive",
+};
diff --git a/src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/components/bot-message.tsx b/src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/components/bot-message.tsx
index 4f03f6b248..50523dae03 100644
--- a/src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/components/bot-message.tsx
+++ b/src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/components/bot-message.tsx
@@ -128,6 +128,51 @@ export const BotMessage = memo(
         ? persistedDuration
         : liveDisplayTime;
 
+    // A message with token usage should still surface the MessageMetadata
+    // pill even when no duration was recorded (e.g. v2 runs that didn't
+    // emit ``build_duration``, historical messages restored from DB).
+    // Without this the user would never see "X tokens" for those.
+    const totalTokens = chat.properties?.usage?.total_tokens;
+    const hasUsage = typeof totalTokens === "number" && totalTokens > 0;
+    const showMetadata = displayTime > 0 || hasUsage;
+
+    // The renderer is data-driven, but content_blocks shows up in two
+    // shapes:
+    //   - Legacy: an "Agent Steps" group wraps the tool calls (and the
+    //     legacy agent also appends a flat TextContent at the top that
+    //     duplicates Message.text). Render the group via the accordion
+    //     and let CustomMarkdownField paint Message.text below — the
+    //     historical "tools on top, text after" layout.
+    //   - Interleaved (post agent-events rewiring): no group, just flat
+    //     tool_use / citation / text items in producer order. Trust the
+    //     content_blocks order and suppress the bubble body so text
+    //     doesn't double-paint.
+    // The signal is: a flat non-text block (tool_use, citation, …) with
+    // no group present means the producer is making an ordering claim.
+    const contentBlocks = chat.content_blocks ?? [];
+    const hasGroup = contentBlocks.some((block) => block.type === "group");
+    const hasFlatNonText = contentBlocks.some(
+      (block) => block.type !== "group" && block.type !== "text",
+    );
+    const hasTextBlock = contentBlocks.some((block) => block.type === "text");
+    const useContentBlockOrdering = !hasGroup && hasFlatNonText;
+    // Suppress the bubble body only when the content blocks actually carry
+    // the answer text. If a producer emits only flat tool_use / citation
+    // items (no TextContent) and stuffs the answer into Message.text,
+    // keep the bubble body so the assistant text isn't hidden.
+    const showBubbleBody =
+      !useContentBlockOrdering || editMessage || !hasTextBlock;
+    // In legacy / pure-text mode, strip a top-level TextContent only when it
+    // duplicates Message.text — those items would render above the grouped
+    // accordion. A divergent text block (text !== Message.text) is kept so
+    // it isn't silently dropped.
+    const displayedContentBlocks = useContentBlockOrdering
+      ? contentBlocks
+      : contentBlocks.filter(
+          (block) =>
+            block.type !== "text" || block.text !== chat.message?.toString(),
+        );
+
     return (
       <>
         

@@ -169,13 +214,15 @@ export const BotMessage = memo(
                       
                         {t("chat.runningStatus")} {formatSeconds(displayTime)}
                       
-                    ) : !thinkingActive && displayTime > 0 ? (
+                    ) : !thinkingActive && showMetadata ? (
                       <>
-                        
-                          {t("chat.finishedIn")}
-                        
+                        {displayTime > 0 && (
+                          
+                            {t("chat.finishedIn")}
+                          
+                        )}
                          0 ? displayTime : undefined}
                           usage={chat.properties?.usage ?? undefined}
                           timestamp={chat.timestamp}
                         />
@@ -184,11 +231,11 @@ export const BotMessage = memo(
                   
                 

 
-                {((chat.content_blocks && chat.content_blocks.length > 0) ||
+                {(displayedContentBlocks.length > 0 ||
                   (isBuilding && lastMessage)) && (
                   
                 )}
 
-                

-                  

-                    

-                      

+                {showBubbleBody && (
+                  

+                    

+                      

                         

-                          {(chatMessage === "" || (isEmpty && !isStreaming)) &&
-                          isBuilding &&
-                          lastMessage ? (
-                            
-                          ) : (
-                            

-                              {editMessage ? (
-                                 setEditMessage(false)}
-                                />
-                              ) : (
-                                <>
-                                  
+                            {(chatMessage === "" ||
+                              (isEmpty && !isStreaming)) &&
+                            isBuilding &&
+                            lastMessage ? (
+                              
+                            ) : (
+                              

+                                {editMessage ? (
+                                   setEditMessage(false)}
                                   />
-                                
-                              )}
-                            

-                          )}
+                                ) : (
+                                  <>
+                                    
+                                  
+                                )}
+                              

+                            )}
+                          

                         

                       

                     

                   

-                

+                )}
               

             

 
diff --git a/src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/components/chat-message.tsx b/src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/components/chat-message.tsx
index 3a2f1bb469..1611056630 100644
--- a/src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/components/chat-message.tsx
+++ b/src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/components/chat-message.tsx
@@ -1,5 +1,6 @@
 import { useEffect, useState } from "react";
 import useFlowStore from "@/stores/flowStore";
+import { isGroupedBlock } from "@/types/chat";
 import type { chatMessagePropsType } from "@/types/components";
 import { BotMessage } from "./bot-message";
 import { ErrorView } from "./error-message";
@@ -27,7 +28,7 @@ export default function ChatMessage({
 
   // Error messages
   if (chat.category === "error") {
-    const blocks = chat.content_blocks ?? [];
+    const blocks = (chat.content_blocks ?? []).filter(isGroupedBlock);
     return (
        {
     expect(updated).toHaveLength(1);
     expect(updated![0].content_blocks).toBeDefined();
     expect(updated![0].content_blocks).toHaveLength(1);
-    expect(updated![0].content_blocks![0].contents[0]).toMatchObject({
+    expect(updated![0].content_blocks![0].contents![0]).toMatchObject({
       type: "tool_use",
       name: "fetch_pokemon",
     });
diff --git a/src/frontend/src/components/core/sanitizedMarkdown/index.tsx b/src/frontend/src/components/core/sanitizedMarkdown/index.tsx
index b70dabc89c..7b220066c1 100644
--- a/src/frontend/src/components/core/sanitizedMarkdown/index.tsx
+++ b/src/frontend/src/components/core/sanitizedMarkdown/index.tsx
@@ -109,7 +109,7 @@ export const SanitizedMarkdown = ({
               );
             },
             // biome-ignore lint/suspicious/noExplicitAny: pre-existing react-markdown override; proper type cleanup out of scope for this PR
-            code: ({ node, inline, className, children, ...props }: any) => {
+            code: ({ node, className, children, ...props }: any) => {
               let content = children as string;
               if (
                 Array.isArray(children) &&
@@ -130,9 +130,18 @@ export const SanitizedMarkdown = ({
                   }
                 }
 
+                // react-markdown v8+ removed the `inline` prop. Detect a
+                // fenced code block by the presence of either a
+                // `language-X` className (hinted block) or multi-line
+                // content (unhinted block). Anything else is inline
+                // backtick code — rendering it as a `
`/`
`
+                // CodeTabsComponent inside a markdown `
` triggers the
+                // 'div cannot be a descendant of p' hydration error.
                 const match = /language-(\w+)/.exec(className || "");
+                const isBlock =
+                  match !== null || (content?.includes("\n") ?? false);
 
-                return !inline ? (
+                return isBlock ? (
                    {
       ]);
     });
   });
+
+  describe("build_duration stamping on output-node success", () => {
+    // The AG-UI translator drops the backend's per-vertex duration, so the
+    // bridge has to derive ``build_duration`` client-side and stamp it onto
+    // the last bot message. These pin the gates (output-type only, success
+    // only, buildStartTime required) so a regression would surface here
+    // instead of in a playwright run.
+
+    const OUTPUT_NODE_ID = "ChatOutput-XYZ";
+    const NON_OUTPUT_NODE_ID = "LanguageModelComponent-ABC";
+
+    const seedFlow = (nodes: Array<{ id: string; type: string }>) => {
+      useFlowStore.setState({
+        nodes: nodes.map((n) => ({
+          id: n.id,
+          // Only ``data.type`` matters here; the rest of the node shape is
+          // surplus for the helper but required by AllNodeType typing.
+          data: { type: n.type },
+        })) as unknown as ReturnType["nodes"],
+      });
+    };
+
+    const seedBotMessage = (msg: {
+      id: string;
+      properties?: Record;
+    }) => {
+      // Stamping reads ``useMessagesStore`` when React Query returns nothing
+      // (the shareable-playground / IOModal path), so seed it there.
+      const { useMessagesStore } = require("@/stores/messagesStore");
+      useMessagesStore.setState({
+        messages: [
+          {
+            id: msg.id,
+            sender: "Machine",
+            sender_name: "AI",
+            session_id: "s1",
+            flow_id: "f1",
+            text: "Hi",
+            files: [],
+            timestamp: "2026-05-28T00:00:00Z",
+            properties: msg.properties ?? {},
+            content_blocks: [],
+            category: "message",
+          },
+        ],
+      });
+    };
+
+    beforeEach(() => {
+      const { useMessagesStore } = require("@/stores/messagesStore");
+      const { queryClient } = require("@/contexts");
+      useMessagesStore.setState({ messages: [] });
+      useFlowStore.setState({ nodes: [], buildStartTime: null });
+      // Clear the React Query messages cache so the scoping test below
+      // doesn't leak seeded sessions into the Zustand-fallback tests.
+      queryClient.clear();
+    });
+
+    it("stamps build_duration on the last bot message when an output node finishes successfully", () => {
+      seedFlow([{ id: OUTPUT_NODE_ID, type: "ChatOutput" }]);
+      seedBotMessage({ id: "m1" });
+      useFlowStore.setState({ buildStartTime: Date.now() - 1500 });
+
+      applyStateDelta(
+        [
+          {
+            op: "add",
+            path: `/nodes/${OUTPUT_NODE_ID}`,
+            value: { status: "success", output: { results: {} } },
+          },
+        ],
+        "run-1",
+        new Set(),
+      );
+
+      const { useMessagesStore } = require("@/stores/messagesStore");
+      const stamped = useMessagesStore.getState().messages[0];
+      expect(stamped.properties.build_duration).toBeGreaterThanOrEqual(1500);
+      // Resets so the next segment is measured fresh.
+      expect(useFlowStore.getState().buildStartTime).not.toBeNull();
+    });
+
+    it("does not stamp when the finishing node is not an output type", () => {
+      seedFlow([{ id: NON_OUTPUT_NODE_ID, type: "LanguageModelComponent" }]);
+      seedBotMessage({ id: "m1" });
+      useFlowStore.setState({ buildStartTime: Date.now() - 500 });
+
+      applyStateDelta(
+        [
+          {
+            op: "add",
+            path: `/nodes/${NON_OUTPUT_NODE_ID}`,
+            value: { status: "success", output: { results: {} } },
+          },
+        ],
+        "run-1",
+        new Set(),
+      );
+
+      const { useMessagesStore } = require("@/stores/messagesStore");
+      const msg = useMessagesStore.getState().messages[0];
+      expect(msg.properties.build_duration).toBeUndefined();
+    });
+
+    it("does not stamp when buildStartTime is missing", () => {
+      seedFlow([{ id: OUTPUT_NODE_ID, type: "ChatOutput" }]);
+      seedBotMessage({ id: "m1" });
+      useFlowStore.setState({ buildStartTime: null });
+
+      applyStateDelta(
+        [
+          {
+            op: "add",
+            path: `/nodes/${OUTPUT_NODE_ID}`,
+            value: { status: "success", output: { results: {} } },
+          },
+        ],
+        "run-1",
+        new Set(),
+      );
+
+      const { useMessagesStore } = require("@/stores/messagesStore");
+      const msg = useMessagesStore.getState().messages[0];
+      expect(msg.properties.build_duration).toBeUndefined();
+    });
+
+    it("does not overwrite a build_duration already on the message (nested-segment guard)", () => {
+      seedFlow([{ id: OUTPUT_NODE_ID, type: "ChatOutput" }]);
+      seedBotMessage({ id: "m1", properties: { build_duration: 999 } });
+      useFlowStore.setState({ buildStartTime: Date.now() - 5000 });
+
+      applyStateDelta(
+        [
+          {
+            op: "add",
+            path: `/nodes/${OUTPUT_NODE_ID}`,
+            value: { status: "success", output: { results: {} } },
+          },
+        ],
+        "run-1",
+        new Set(),
+      );
+
+      const { useMessagesStore } = require("@/stores/messagesStore");
+      const msg = useMessagesStore.getState().messages[0];
+      expect(msg.properties.build_duration).toBe(999);
+    });
+
+    it("scopes the stamp to the running session, not another session's cache", () => {
+      // D1 regression: the bot-message lookup must be scoped to the running
+      // flow/session. The unscoped form walked every messages cache and could
+      // stamp build_duration onto a bot message from a different session the
+      // user had open in the same tab.
+      const { queryClient } = require("@/contexts");
+      const MESSAGES_QUERY_KEY = "useGetMessagesQuery";
+      const botMsg = (id: string) => ({
+        id,
+        sender: "Machine",
+        sender_name: "AI",
+        text: "hi",
+        files: [],
+        timestamp: "2026-05-28T00:00:00Z",
+        properties: {} as Record,
+        content_blocks: [],
+        category: "message",
+      });
+      const otherKey = [
+        MESSAGES_QUERY_KEY,
+        { id: "f1", session_id: "other-session" },
+      ];
+      const runningKey = [
+        MESSAGES_QUERY_KEY,
+        { id: "f1", session_id: "running-session" },
+      ];
+      queryClient.setQueryData(otherKey, [botMsg("other-msg")]);
+      queryClient.setQueryData(runningKey, [botMsg("running-msg")]);
+
+      seedFlow([{ id: OUTPUT_NODE_ID, type: "ChatOutput" }]);
+      useFlowStore.setState({ buildStartTime: Date.now() - 1500 });
+
+      applyStateDelta(
+        [
+          {
+            op: "add",
+            path: `/nodes/${OUTPUT_NODE_ID}`,
+            value: { status: "success", output: { results: {} } },
+          },
+        ],
+        "run-1",
+        new Set(),
+        "f1",
+        "running-session",
+      );
+
+      const running = queryClient.getQueryData(runningKey) as Array<{
+        properties: { build_duration?: number };
+      }>;
+      const other = queryClient.getQueryData(otherKey) as Array<{
+        properties: { build_duration?: number };
+      }>;
+      expect(running[0].properties.build_duration).toBeGreaterThanOrEqual(1500);
+      expect(other[0].properties.build_duration).toBeUndefined();
+    });
+  });
 });
diff --git a/src/frontend/src/controllers/API/agui/run-flow-bridge.ts b/src/frontend/src/controllers/API/agui/run-flow-bridge.ts
index 74ed69b1f9..afb79df066 100644
--- a/src/frontend/src/controllers/API/agui/run-flow-bridge.ts
+++ b/src/frontend/src/controllers/API/agui/run-flow-bridge.ts
@@ -10,15 +10,22 @@
 
 import { type BaseEvent, EventType } from "@ag-ui/client";
 import { handleMessageEvent } from "@/components/core/playgroundComponent/chat-view/utils/message-event-handler";
+import {
+  findLastBotMessage,
+  updateMessageProperties,
+} from "@/components/core/playgroundComponent/chat-view/utils/message-utils";
 import { BuildStatus } from "@/constants/enums";
+import { persistMessageProperties } from "@/controllers/API/helpers/persist-message-properties";
 import useAlertStore from "@/stores/alertStore";
 import useFlowStore from "@/stores/flowStore";
+import { useMessagesStore } from "@/stores/messagesStore";
 import type {
   ChatInputType,
   ChatOutputType,
   VertexBuildTypeAPI,
   VertexDataTypeAPI,
 } from "@/types/api";
+import { isOutputType } from "@/utils/reactflowUtils";
 import {
   buildWorkflowRunRequest,
   createWorkflowAgent,
@@ -112,6 +119,8 @@ export function applyStateDelta(
   ops: JsonPatchOp[],
   runId: string,
   nodeIds: Set,
+  flowId?: string,
+  sessionId?: string,
 ): void {
   const flowStore = useFlowStore.getState();
   for (const op of ops) {
@@ -155,10 +164,99 @@ export function applyStateDelta(
         artifacts: null,
       };
       flowStore.addDataToFlowPool(entry, nodeId);
+
+      // When a ChatOutput (or other output type) finishes, stamp the
+      // elapsed segment duration onto the bot message it just produced.
+      // Mirrors the v1 ``end_vertex`` callback in ``buildUtils.ts`` so
+      // the playground header can render the "Finished in" pill and
+      // MessageMetadata can show its duration. The AG-UI translator
+      // does not propagate ``build_duration`` from the backend, so this
+      // is the only path that sets it for v2 runs.
+      if (value.status === "success") {
+        stampSegmentDurationForOutputNode(nodeId, flowId, sessionId);
+      }
     }
   }
 }
 
+/**
+ * If ``nodeId`` is an output-type node, compute the time since the
+ * flow's ``buildStartTime`` and persist it as ``build_duration`` on the
+ * last bot message in the matching React Query cache (and the Zustand
+ * fallback used by the shareable playground). ``flowId``/``sessionId`` scope
+ * the lookup to the running session so the duration can't land on a bot
+ * message from a different session's cache. Resets ``buildStartTime`` so the
+ * next segment is measured fresh.
+ *
+ * Mirrors the per-vertex segment logic in ``buildUtils.ts`` so the v2
+ * AG-UI bridge produces the same on-message metadata the v1 build
+ * callbacks do.
+ */
+function stampSegmentDurationForOutputNode(
+  nodeId: string,
+  flowId?: string,
+  sessionId?: string,
+): void {
+  const flowState = useFlowStore.getState();
+  const node = flowState.nodes.find((n) => n.id === nodeId);
+  const nodeType = node?.data?.type as string | undefined;
+  if (!nodeType || !isOutputType(nodeType) || !flowState.buildStartTime) {
+    return;
+  }
+  const segmentDurationMs = Date.now() - flowState.buildStartTime;
+
+  const found = findLastBotMessage(flowId, sessionId);
+  if (found && !found.message.properties?.build_duration) {
+    updateMessageProperties(found.message.id!, found.queryKey, {
+      build_duration: segmentDurationMs,
+    });
+
+    const storeMsg = useMessagesStore
+      .getState()
+      .messages.find((m) => m.id === found.message.id);
+    if (storeMsg) {
+      useMessagesStore.getState().updateMessage({
+        ...storeMsg,
+        properties: {
+          ...storeMsg.properties,
+          build_duration: segmentDurationMs,
+        },
+      });
+    }
+
+    persistMessageProperties(found.message.id!, {
+      ...found.message,
+      properties: {
+        ...found.message.properties,
+        build_duration: segmentDurationMs,
+      },
+    });
+  } else if (!found) {
+    // Shareable-playground fallback: React Query cache is empty, so look
+    // for the last bot message in the Zustand store and stamp there.
+    const storeMessages = useMessagesStore.getState().messages;
+    for (let i = storeMessages.length - 1; i >= 0; i--) {
+      const msg = storeMessages[i];
+      if (msg.sender === "Machine" && !msg.properties?.build_duration) {
+        const updatedProperties = {
+          ...msg.properties,
+          build_duration: segmentDurationMs,
+        };
+        useMessagesStore.getState().updateMessage({
+          ...msg,
+          properties: updatedProperties,
+        });
+        if (msg.id) {
+          persistMessageProperties(msg.id, { properties: updatedProperties });
+        }
+        break;
+      }
+    }
+  }
+
+  flowState.setBuildStartTime(Date.now());
+}
+
 /**
  * Run a workflow through the AG-UI service and update flowStore as events
  * arrive. Resolves when the run ends (RUN_FINISHED or RUN_ERROR), the
@@ -203,6 +301,11 @@ export async function runFlowAGUI(
   const flowStore = useFlowStore.getState();
   const setErrorData = useAlertStore.getState().setErrorData;
   const touchedNodeIds = new Set();
+  // ``setIsBuilding(true)`` (called upstream by ``buildFlow``) clears
+  // ``buildStartTime``. Initialise it here so the segment duration the
+  // per-vertex success handler stamps onto bot messages is measured
+  // from the actual moment the run started, not from a stale tick.
+  flowStore.setBuildStartTime(Date.now());
   // Server derives run_id from session_id/flow_id and announces it via
   // RUN_STARTED. Until that arrives we have no id to stamp on flow-pool
   // entries, so we fall back to an empty string (matches the legacy
@@ -237,7 +340,8 @@ export async function runFlowAGUI(
     setRunId: (r) => {
       runId = r;
     },
-    applyDelta: (ops) => applyStateDelta(ops, runId, touchedNodeIds),
+    applyDelta: (ops) =>
+      applyStateDelta(ops, runId, touchedNodeIds, opts.flowId, opts.threadId),
     handleCustomEvent: (eventType, data) => handleMessageEvent(eventType, data),
     onFinished: () => {
       terminalEventSeen = true;
diff --git a/src/frontend/src/modals/IOModal/components/chatView/chatMessage/chat-message.tsx b/src/frontend/src/modals/IOModal/components/chatView/chatMessage/chat-message.tsx
index 1d9e1ae082..b9e58dd87c 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/chatMessage/chat-message.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/chatMessage/chat-message.tsx
@@ -9,6 +9,7 @@ import { CustomProfileIcon } from "@/customization/components/custom-profile-ico
 import { ENABLE_DATASTAX_LANGFLOW } from "@/customization/feature-flags";
 import useFlowStore from "@/stores/flowStore";
 import useFlowsManagerStore from "@/stores/flowsManagerStore";
+import { isGroupedBlock } from "@/types/chat";
 import Robot from "../../../../../assets/robot.png";
 import IconComponent, {
   ForwardedIconComponent,
@@ -130,6 +131,43 @@ export default function ChatMessage({
   const isEmpty = decodedMessage?.trim() === "";
   const { mutate: updateMessageMutation } = useUpdateMessage();
 
+  // The renderer is data-driven, but content_blocks shows up in two
+  // shapes:
+  //   - Legacy: an "Agent Steps" group wraps the tool calls (and the
+  //     legacy agent also appends a flat TextContent at the top that
+  //     duplicates Message.text). Render the group via the accordion
+  //     and let CustomMarkdownField paint Message.text below — the
+  //     historical "tools on top, text after" layout.
+  //   - Interleaved (post agent-events rewiring): no group, just flat
+  //     tool_use / citation / text items in producer order. Trust the
+  //     content_blocks order and suppress the bubble body so text
+  //     doesn't double-paint.
+  // The signal is: a flat non-text block (tool_use, citation, …) with
+  // no group present means the producer is making an ordering claim.
+  const contentBlocks = chat.content_blocks ?? [];
+  const hasGroup = contentBlocks.some((block) => block.type === "group");
+  const hasFlatNonText = contentBlocks.some(
+    (block) => block.type !== "group" && block.type !== "text",
+  );
+  const hasTextBlock = contentBlocks.some((block) => block.type === "text");
+  const useContentBlockOrdering = !hasGroup && hasFlatNonText;
+  // Suppress the bubble body only when the content blocks actually carry
+  // the answer text. If a producer emits only flat tool_use / citation
+  // items (no TextContent) and stuffs the answer into Message.text, keep
+  // the bubble body so the assistant text isn't hidden.
+  const showBubbleBody =
+    !useContentBlockOrdering || editMessage || !hasTextBlock;
+  // In legacy / pure-text mode, strip a top-level TextContent only when it
+  // duplicates Message.text — those items would render above the grouped
+  // accordion. A divergent text block (text !== Message.text) is kept so it
+  // isn't silently dropped.
+  const displayedContentBlocks = useContentBlockOrdering
+    ? contentBlocks
+    : contentBlocks.filter(
+        (block) =>
+          block.type !== "text" || block.text !== chat.message?.toString(),
+      );
+
   const handleEditMessage = (message: string) => {
     updateMessageMutation(
       {
@@ -191,7 +229,10 @@ export default function ChatMessage({
   ) : null;
 
   if (chat.category === "error") {
-    const blocks = chat.content_blocks ?? [];
+    // ErrorView walks block.contents on each entry, so filter to the grouped
+    // ContentBlock shape (it has `contents`). Flat ContentType items can now
+    // also appear in content_blocks and would otherwise crash ErrorView.
+    const blocks = (chat.content_blocks ?? []).filter(isGroupedBlock);
 
     return (
       
             

-            {chat.content_blocks && chat.content_blocks.length > 0 && (
+            {displayedContentBlocks.length > 0 && (
               
             )}
-            {!chat.isSend ? (
+            {!chat.isSend && !showBubbleBody ? null : !chat.isSend ? (
               

                 

                   {hidden && chat.thought && chat.thought !== "" && (
diff --git a/src/frontend/src/types/chat/__tests__/isGroupedBlock.test.ts b/src/frontend/src/types/chat/__tests__/isGroupedBlock.test.ts
new file mode 100644
index 0000000000..7ef897a863
--- /dev/null
+++ b/src/frontend/src/types/chat/__tests__/isGroupedBlock.test.ts
@@ -0,0 +1,78 @@
+import { type ContentBlockItem, isGroupedBlock } from "..";
+
+// Backend's BaseContent serializes `contents: []` by default on every leaf,
+// so the structural fallback in `isGroupedBlock` MUST NOT fire when the
+// discriminator says the item is a flat ContentType — otherwise a flat
+// CodeContent or CitationContent with a `title` field gets misrouted into
+// the grouped accordion and silently disappears from the UI.
+
+describe("isGroupedBlock", () => {
+  it("returns true when type === 'group' (discriminator wins)", () => {
+    const block: ContentBlockItem = {
+      type: "group",
+      title: "Agent Steps",
+      contents: [],
+    };
+    expect(isGroupedBlock(block)).toBe(true);
+  });
+
+  it("returns false for a flat CodeContent that happens to have a title", () => {
+    // Backend serializes optional `title` and inherited `contents: []` on
+    // every CodeContent, so the structural fallback would misclassify this
+    // as a group if it didn't gate on the discriminator first.
+    const code = {
+      type: "code",
+      title: "snippet.py",
+      code: "print('hi')",
+      language: "python",
+      contents: [],
+    } as unknown as ContentBlockItem;
+    expect(isGroupedBlock(code)).toBe(false);
+  });
+
+  it("returns false for a flat CitationContent with title + empty contents", () => {
+    const citation = {
+      type: "citation",
+      title: "Source",
+      url: "https://example.com",
+      contents: [],
+    } as unknown as ContentBlockItem;
+    expect(isGroupedBlock(citation)).toBe(false);
+  });
+
+  it("returns false for a flat ToolContent (no title, has contents)", () => {
+    const tool = {
+      type: "tool_use",
+      name: "fetch",
+      contents: [],
+    } as unknown as ContentBlockItem;
+    expect(isGroupedBlock(tool)).toBe(false);
+  });
+
+  it("returns true for a legacy ContentBlock dict that lacks a type field", () => {
+    // Pre-discriminator persisted shape: only title + contents, no `type`.
+    // This is the case the structural fallback exists for.
+    const legacy = {
+      title: "Agent Steps",
+      contents: [{ type: "tool_use", name: "x" }],
+      allow_markdown: true,
+    } as unknown as ContentBlockItem;
+    expect(isGroupedBlock(legacy)).toBe(true);
+  });
+
+  it("returns false for a legacy-shaped dict that lacks a contents array", () => {
+    // Without contents we can't tell it's a group; treat as flat.
+    const malformed = {
+      title: "Stray title",
+    } as unknown as ContentBlockItem;
+    expect(isGroupedBlock(malformed)).toBe(false);
+  });
+
+  it("returns false for a TextContent", () => {
+    const text = {
+      type: "text",
+      text: "hello",
+    } as unknown as ContentBlockItem;
+    expect(isGroupedBlock(text)).toBe(false);
+  });
+});
diff --git a/src/frontend/src/types/chat/index.ts b/src/frontend/src/types/chat/index.ts
index 81a88e5540..0eab19a6a0 100644
--- a/src/frontend/src/types/chat/index.ts
+++ b/src/frontend/src/types/chat/index.ts
@@ -25,7 +25,7 @@ export type ChatMessageType = {
   icon?: string;
   category?: string;
   properties?: PropertiesType;
-  content_blocks?: ContentBlock[];
+  content_blocks?: ContentBlockItem[];
 };
 
 export type SourceType = {
@@ -85,11 +85,20 @@ export type FlowPoolObjectType = {
 // Base content type
 export interface BaseContent {
   type: string;
+  // Optional stable identity carried across re-emissions of the same logical
+  // block. Set by producers that have a natural id (e.g. LangChain
+  // tool_call_id). Consumers fall back to position-derived dedup when absent.
+  id?: string;
   duration?: number;
   header?: {
     title?: string;
     icon?: string;
   };
+  // Nested content. Leaf types leave this empty; container-shaped types
+  // (ContentBlock, multimodal ToolContent, multi-step ReasoningContent)
+  // populate it. Typed as ContentBlockItem so nested ContentBlock groups
+  // are allowed, matching the backend's discriminated union.
+  contents?: ContentBlockItem[];
 }
 
 // Individual content types
@@ -128,11 +137,69 @@ export interface CodeContent extends BaseContent {
 export interface ToolContent extends BaseContent {
   type: "tool_use";
   name?: string;
-  tool_input: Record;
+  // The backend serializes this field as `tool_input` by default but
+  // emits it under the `input` alias when AG-UI / any other path runs
+  // model_dump with by_alias=True. Accept both shapes; renderers should
+  // prefer `tool_input` and fall back to `input`.
+  tool_input?: Record;
+  input?: Record;
   output?: JSONValue;
   error?: JSONValue | string;
 }
 
+export interface ImageContent extends BaseContent {
+  type: "image";
+  urls?: string[];
+  base64?: string;
+  mime_type?: string;
+  caption?: string;
+}
+
+export interface AudioContent extends BaseContent {
+  type: "audio";
+  urls?: string[];
+  base64?: string;
+  mime_type?: string;
+  duration?: number;
+  transcript?: string;
+}
+
+export interface VideoContent extends BaseContent {
+  type: "video";
+  urls?: string[];
+  base64?: string;
+  mime_type?: string;
+  duration?: number;
+}
+
+export interface FileContent extends BaseContent {
+  type: "file";
+  urls?: string[];
+  mime_type?: string;
+  filename?: string;
+}
+
+export interface ReasoningContent extends BaseContent {
+  type: "reasoning";
+  text: string;
+}
+
+export interface UsageContent extends BaseContent {
+  type: "usage";
+  input_tokens?: number;
+  output_tokens?: number;
+  model?: string;
+}
+
+export interface CitationContent extends BaseContent {
+  type: "citation";
+  url?: string;
+  title?: string;
+  cited_text?: string;
+  start_index?: number;
+  end_index?: number;
+}
+
 // Union type for all content types
 export type ContentType =
   | ErrorContent
@@ -140,15 +207,50 @@ export type ContentType =
   | MediaContent
   | JSONContent
   | CodeContent
-  | ToolContent;
+  | ToolContent
+  | ImageContent
+  | AudioContent
+  | VideoContent
+  | FileContent
+  | ReasoningContent
+  | UsageContent
+  | CitationContent;
 
-// Updated ContentBlock interface
-export interface ContentBlock {
+// A titled group of nested contents. Matches the backend's ContentBlock,
+// which is a member of the ContentType discriminated union with tag "group".
+// Extends BaseContent so it inherits id/header/duration alongside the
+// group-specific fields.
+export interface ContentBlock extends BaseContent {
+  type: "group";
   title: string;
-  contents: ContentType[];
-  allow_markdown: boolean;
+  contents: ContentBlockItem[];
+  // Optional to match the backend default (True) and to tolerate hand-built
+  // or legacy payloads that don't include the field.
+  allow_markdown?: boolean;
   media_url?: string[];
-  component: string;
+  // Legacy field kept for backwards compatibility with older callers.
+  component?: string;
+}
+
+// A content block item can be either a grouped ContentBlock or a flat ContentType
+export type ContentBlockItem = ContentType | ContentBlock;
+
+// Type guard for grouped ContentBlock items. Prefers the discriminator
+// `type === "group"` (new payloads). Falls back to a structural check ONLY
+// when `type` is absent, so legacy ContentBlock dicts persisted before the
+// discriminator existed (no `type`, only title + contents) are still
+// classified as groups. The fallback must not fire on flat ContentType
+// items that happen to carry a `title` (CodeContent, CitationContent) and
+// an empty inherited `contents: []` -- backend BaseContent serializes
+// contents on every item, so the type check is what keeps them out.
+export function isGroupedBlock(item: ContentBlockItem): item is ContentBlock {
+  if (item.type === "group") return true;
+  if (item.type) return false;
+  return (
+    "title" in item &&
+    "contents" in item &&
+    Array.isArray((item as { contents?: unknown }).contents)
+  );
 }
 
 export interface PlaygroundEvent {
@@ -158,7 +260,7 @@ export interface PlaygroundEvent {
   allow_markdown?: boolean;
   icon?: string | null;
   sender_name: string;
-  content_blocks?: ContentBlock[] | null;
+  content_blocks?: ContentBlockItem[] | null;
   files?: string[];
   text?: string;
   timestamp?: string;
diff --git a/src/frontend/src/types/messages/index.ts b/src/frontend/src/types/messages/index.ts
index 56c065ee1b..dcaed93aa2 100644
--- a/src/frontend/src/types/messages/index.ts
+++ b/src/frontend/src/types/messages/index.ts
@@ -1,4 +1,4 @@
-import type { ContentBlock, UsageType } from "../chat";
+import type { ContentBlockItem, UsageType } from "../chat";
 
 type Message = {
   flow_id: string;
@@ -31,7 +31,7 @@ type Message = {
     usage?: UsageType | null;
     [key: string]: unknown;
   };
-  content_blocks?: ContentBlock[];
+  content_blocks?: ContentBlockItem[];
 };
 
 export type { Message };
diff --git a/src/lfx/src/lfx/schema/content_types.py b/src/lfx/src/lfx/schema/content_types.py
index 2f923453d6..d74c3dfc3d 100644
--- a/src/lfx/src/lfx/schema/content_types.py
+++ b/src/lfx/src/lfx/schema/content_types.py
@@ -75,6 +75,19 @@ class BaseContent(BaseModel):
         ),
     )
 
+    def __init__(self, **data: Any) -> None:
+        super().__init__(**data)
+        # Mark only the discriminator as "set" so partial-update callers using
+        # ``model_dump(exclude_unset=True)`` (notably ``aupdate_messages``)
+        # keep the ``type`` tag. Without it, ``TextContent(text="x")`` would
+        # dump to ``{"text": "x"}`` and the next read-back through
+        # ``MessageRead``'s discriminated union would fail with
+        # ``union_tag_not_found``. Other defaulted fields stay unset so true
+        # exclude_unset semantics survive: a patch like ``ContentBlock(title=
+        # "...")`` no longer overwrites an existing block's ``duration`` /
+        # ``header`` / ``contents`` with their defaults on merge.
+        self.model_fields_set.add("type")
+
     def to_dict(self) -> dict[str, Any]:
         return self.model_dump()
 
@@ -260,17 +273,10 @@ class ContentBlock(BaseContent):
     allow_markdown: bool = Field(default=True)
     media_url: list[str] | None = None
 
-    def __init__(self, **data) -> None:
-        super().__init__(**data)
-        # Mark only the discriminator as "set" so partial-update callers
-        # using ``model_dump(exclude_unset=True)`` (notably
-        # ``aupdate_messages``) still carry the ``type="group"`` field
-        # downstream. Other defaulted fields stay unset so true
-        # exclude_unset semantics survive: a patch like
-        # ``ContentBlock(title="...")`` no longer overwrites an existing
-        # block's ``duration`` / ``header`` / ``contents`` with their
-        # defaults on merge.
-        self.model_fields_set.add("type")
+    # ``__init__`` is inherited from ``BaseContent``: it marks only the
+    # ``type`` discriminator as set, so ``model_dump(exclude_unset=True)``
+    # preserves ``type="group"`` while leaving ``allow_markdown`` and the
+    # other defaulted fields out of partial-update merges.
 
     @field_validator("contents", mode="before")
     @classmethod
diff --git a/src/lfx/src/lfx/schema/message.py b/src/lfx/src/lfx/schema/message.py
index f6df1aef01..a6964f9711 100644
--- a/src/lfx/src/lfx/schema/message.py
+++ b/src/lfx/src/lfx/schema/message.py
@@ -743,7 +743,15 @@ class MessageResponse(DefaultModel):
 
     properties: Properties | None = None
     category: str | None = None
-    content_blocks: list[ContentType | ContentBlock] | None = None
+    # ContentBlock joined the ContentType discriminated union as tag "group",
+    # so list[ContentType] alone covers both flat ContentType (text, image,
+    # tool_use, ...) and the wrapping ContentBlock shape. The previous
+    # `list[ContentType | ContentBlock]` union confused Pydantic's
+    # discriminator: any dict with a `contents` field (which every backend
+    # BaseContent now serializes by default, even as []) was routed to
+    # ContentBlock and rejected with `type: literal_error` for non-group
+    # types like text.
+    content_blocks: list[ContentType] | None = None
     session_metadata: dict | None = None
 
     @field_validator("content_blocks", mode="before")
@@ -754,10 +762,11 @@ class MessageResponse(DefaultModel):
         if isinstance(v, list):
             return [cls.validate_content_blocks(block) for block in v]
         if isinstance(v, dict):
-            # Flat ContentType (has "type" but no "contents") vs grouped ContentBlock
-            if "type" in v and "contents" not in v:
-                return _CONTENT_TYPE_ADAPTER.validate_python(v)
-            return ContentBlock.model_validate(v)
+            # The discriminator on `type` handles flat ContentType vs grouped
+            # ContentBlock uniformly. The old "type and not contents"
+            # heuristic broke once BaseContent started serializing
+            # `contents: []` on every leaf.
+            return _CONTENT_TYPE_ADAPTER.validate_python(v)
         return v
 
     @field_validator("properties", mode="before")
diff --git a/src/lfx/tests/unit/schema/test_message_content_blocks.py b/src/lfx/tests/unit/schema/test_message_content_blocks.py
index 20d28c2a16..2d6fe387b3 100644
--- a/src/lfx/tests/unit/schema/test_message_content_blocks.py
+++ b/src/lfx/tests/unit/schema/test_message_content_blocks.py
@@ -18,7 +18,7 @@ from lfx.schema.content_types import (
     ToolContent,
 )
 from lfx.schema.data import Data
-from lfx.schema.message import ErrorMessage, Message
+from lfx.schema.message import ErrorMessage, Message, MessageResponse
 from lfx.schema.properties import Source
 from lfx.utils.constants import (
     MESSAGE_SENDER_AI,
@@ -544,3 +544,149 @@ class TestContentBlockExcludeUnset:
         # ``header``, ``media_url`` must stay out of the dump so they
         # don't clobber existing values on merge.
         assert dump.keys() <= {"type", "title", "allow_markdown"}
+
+
+class TestMessageResponseContentBlocksValidation:
+    """Regression: MessageResponse must accept every ContentType in content_blocks.
+
+    The agent now appends flat items (TextContent, ToolContent, etc.) to
+    ``content_blocks`` instead of nesting them inside a ContentBlock
+    group. The serialized form lands on the API as
+    ``{"type": "text", ...}`` etc. If MessageResponse only validates
+    ContentBlock, the API errors out with
+    ``content_blocks.type: Input should be 'group'`` and the entire
+    chat response fails to render.
+    """
+
+    _BASE_KWARGS = {
+        "sender": "Machine",
+        "sender_name": "AI",
+        "session_id": "s",
+        "text": "Doing well, thanks!",
+        "edit": False,
+    }
+
+    def test_flat_text_content_validates(self):
+        # The exact shape the agent's setter appends after `message.text =`.
+        payload = {
+            **self._BASE_KWARGS,
+            "content_blocks": [
+                {
+                    "type": "text",
+                    "id": None,
+                    "duration": None,
+                    "header": {},
+                    "contents": [],
+                    "text": "Doing well, thanks!",
+                },
+            ],
+        }
+        resp = MessageResponse.model_validate(payload)
+        assert resp.content_blocks is not None
+        assert len(resp.content_blocks) == 1
+        assert resp.content_blocks[0].type == "text"
+        assert resp.content_blocks[0].text == "Doing well, thanks!"
+
+    def test_flat_tool_content_validates(self):
+        payload = {
+            **self._BASE_KWARGS,
+            "content_blocks": [
+                {
+                    "type": "tool_use",
+                    "name": "search",
+                    "tool_input": {"q": "x"},
+                    "contents": [],
+                },
+            ],
+        }
+        resp = MessageResponse.model_validate(payload)
+        assert resp.content_blocks[0].type == "tool_use"
+        assert resp.content_blocks[0].name == "search"
+
+    def test_mixed_flat_and_grouped_validates(self):
+        # Chronological-events shape: a tool call, then a grouped block
+        # (legacy), then the setter-appended text. All three must round-trip.
+        payload = {
+            **self._BASE_KWARGS,
+            "content_blocks": [
+                {"type": "tool_use", "name": "calc", "contents": []},
+                {
+                    "type": "group",
+                    "title": "Agent Steps",
+                    "contents": [],
+                    "allow_markdown": True,
+                },
+                {"type": "text", "text": "done", "contents": []},
+            ],
+        }
+        resp = MessageResponse.model_validate(payload)
+        types = [b.type for b in resp.content_blocks]
+        assert types == ["tool_use", "group", "text"]
+
+    def test_validates_agent_setter_emitted_content_blocks(self):
+        # End-to-end: the agent's setter path. process_agent_events calls
+        # `message.text = "..."` post-construction and the setter appends
+        # a TextContent to content_blocks. The serialized payload from
+        # that Message must validate as a MessageResponse without the
+        # discriminated-union rejecting the flat TextContent shape.
+        message = Message(
+            text="",
+            sender="Machine",
+            sender_name="AI",
+            session_id="s",
+        )
+        # Trigger the setter; this is what handle_on_chain_end does.
+        message.text = "Doing well, thanks!"
+        assert any(isinstance(b, TextContent) for b in message.content_blocks)
+
+        # Round-trip the model_dump payload, which is what the API sends
+        # over the wire.
+        payload = {**self._BASE_KWARGS, "content_blocks": [b.model_dump() for b in message.content_blocks]}
+        resp = MessageResponse.model_validate(payload)
+        assert resp.content_blocks is not None
+        assert resp.content_blocks[0].type == "text"
+        assert resp.content_blocks[0].text == "Doing well, thanks!"
+
+
+class TestContentTypeExcludeUnsetPreservesDiscriminator:
+    """Regression: ContentType.model_dump(exclude_unset=True) must keep type.
+
+    aupdate_messages applies a partial-update via:
+        msg.sqlmodel_update(message.model_dump(exclude_unset=True, exclude_none=True))
+
+    If TextContent dumps to ``{"text": "..."}`` (stripping the defaulted
+    ``type`` field), the stored row's content_blocks entry has no
+    discriminator. The next read-back through MessageRead's discriminated
+    union fails with ``union_tag_not_found``, aupdate_messages raises
+    ValueError, astore_message silently falls through to aadd_messages,
+    and a duplicate row gets inserted with a new id — which the chat
+    view renders as a duplicate bubble.
+
+    The fix is in BaseContent.__init__: mark every field as set so
+    exclude_unset keeps the discriminator.
+    """
+
+    def test_text_content_keeps_type(self):
+        dump = TextContent(text="hi").model_dump(exclude_unset=True, exclude_none=True)
+        assert dump.get("type") == "text", f"missing discriminator: {dump}"
+        assert dump.get("text") == "hi"
+
+    def test_tool_content_keeps_type(self):
+        dump = ToolContent(name="search").model_dump(exclude_unset=True, exclude_none=True)
+        assert dump.get("type") == "tool_use", f"missing discriminator: {dump}"
+
+    def test_error_content_keeps_type(self):
+        dump = ErrorContent(reason="boom").model_dump(exclude_unset=True, exclude_none=True)
+        assert dump.get("type") == "error", f"missing discriminator: {dump}"
+
+    def test_content_block_group_keeps_type(self):
+        dump = ContentBlock(title="Agent Steps", contents=[]).model_dump(
+            exclude_unset=True,
+            exclude_none=True,
+        )
+        assert dump.get("type") == "group", f"missing discriminator: {dump}"
+        assert dump.get("title") == "Agent Steps"
+
+    def test_image_content_keeps_type(self):
+        dump = ImageContent(urls=["x"]).model_dump(exclude_unset=True, exclude_none=True)
+        assert dump.get("type") == "image", f"missing discriminator: {dump}"