Generative UI elements now working on tool-call while using a custom backend

[BinariusConsulting]

I am having trouble displaying Generative UI components on tool-calls. I have a custom backend that has a langgraph graph with an endpoint exposed for streaming events. As written in documentation, I have created a custom model adaptor to stream and view the messages on assistant-ui. The text streams work well and are displayed correctly but the tool-call streams seem to start rendering for a split second and then disappear.

const MyModelAdapter: ChatModelAdapter = {
  async *run({ messages, abortSignal, context }) {
    console.log("MyModelAdapter run called with:", {
      messagesCount: messages.length,
      lastMessage: messages[messages.length - 1],
      context,
    });

    try {
      const stream = await backendApi({ messages, abortSignal, context });
      const decoder = new TextDecoder();

      let text = "";
      while (true) {
        const { done, value } = await stream.read();
        if (done) {
          console.log("Stream completed");
          break;
        }

        const chunk = decoder.decode(value);

        const lines = chunk.split("\n").filter(Boolean);

        for (const line of lines) {
          try {
            if (line.startsWith("data: ")) {
              const jsonStr = line.slice(6);
              const data = JSON.parse(jsonStr);
              // console.log("Parsed event data:", data);

              text += data.data?.chunk?.kwargs?.content || "";

              const toolCalls = data.data?.chunk?.kwargs?.tool_calls || [];
              const content: ThreadAssistantContentPart[] = [
                { type: "text", text },
              ];

              for (const toolCall of toolCalls) {
                content.push({
                  type: "tool-call",
                  toolName: toolCall.name,
                  args: toolCall.args,
                  argsText: toolCall.argsText,
                  toolCallId: toolCall.toolCallId,
                  result: toolCall.result,
                } as ThreadAssistantContentPart);
              }

              yield { content };
            }
          } catch (e) {
            console.error("Failed to parse chunk:", e, "Line:", line);
          }
        }
      }
    } catch (error) {
      console.error("MyModelAdapter run error:", error);
      throw error;
    }
  },
};

Below is one of my tools

export const UserSnapshotTool = makeAssistantToolUI<
  UserSnapshotToolArgs,
  string
>({
  toolName: "user_snapshot",
  render: function UserSnapshotUI({ args, result }) {
    console.log("UserSnapshotTool render", { args, result });
    let resultObj: UserSnapshotToolResult | { error: string };
    try {
      resultObj = result ? JSON.parse(result) : {};
    } catch (e) {
      resultObj = { error: result! };
    }

    return (
      <div className="mb-4 flex flex-col items-center gap-2">
        <pre className="whitespace-pre-wrap break-all text-center">
          user_snapshot({JSON.stringify(args)})
        </pre>
        {"snapshot" in resultObj && (
          <UserSnapshot
            email={args.email}
            name={(resultObj.snapshot as UserSnapshotToolResult).name}
            phone={(resultObj.snapshot as UserSnapshotToolResult).phone}
            organization={
              (resultObj.snapshot as UserSnapshotToolResult).organization
            }
          />
        )}
        {"error" in resultObj && (
          <p className="text-red-500">{resultObj.error}</p>
        )}
      </div>
    );
  },
});

I see a console log of UserSnapshotTool render but the component itself just starts to render and then disappears. Let me know if you need more information.

[eladlachmi]

I implemented a very similar tool UI using version 0.8.6 and unstyled components, and am having the same issue. I've combed through all of the examples and issues in the repo and couldn't get it to work.

It definitely might be a skill issue on my part :)
Any assistance would be greatly appreciated non-the-less.

[JatSh1804]

@eladlachmi @BinariusConsulting More context will really help, but considering all the info i think the problem's due to the way render function is tryna extract multiple arguments from the streamed response. The UI data should be passed into the args field itself and not in result. You can try schematic: docs

[krankos]

I'm facing the same issue

[krankos]

@JatSh1804 I think what you are referring to is an example of a frontend tool. The tool renders a chart from given data, that's why it doesn't have a result. What @BinariusConsulting is referring to is different.

[zainozzaini]

I'm facing the same issue

[arindambiswas1]

I am doing something similar and can see the custom tool UI in chat, but the follow up questions are crashing fastapi, langraph backend with this error

Error code: 400 - {'error': {'message': "An assistant message with 'tool_calls' must be followed by tool messages responding to each 'tool_call_id'. The following tool_call_ids did not have response messages: call_OeSie4McRg1pc2D578IX9JvB", 'type': 'invalid_request_error', 'param': 'messages.[5].role', 'code': None}}

Here is my Adapter

const MyModelAdapter: ChatModelAdapter = {
  async *run({ messages, abortSignal, context }) {
    console.log("MyModelAdapter run called with:", {
      messagesCount: messages.length,
      lastMessage: messages[messages.length - 1],
      context,
    });
    const parseToolLine = createToolCallParser();
    try {
      console.log(messages)
      const stream = await fetch("http://localhost:8000/api/chat", {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
        },
        body: JSON.stringify({
          messages,
        }),
        signal: abortSignal,
      });
      const decoder = new TextDecoder();
      const reader = stream.body.getReader();
      let buffer = "";
      let text = "";
      let toolCall = null;
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        buffer += decoder.decode(value, { stream: true });

        const lines = buffer.split("\n");
        buffer = lines.pop();

        for (const line of lines) {
          const ptcall = parseToolLine(line);
          if (ptcall) {
            toolCall = ptcall;
          }
          const match = line.match(/0:"(.*)"/);
          if (match && match[1]) {
            const token = match[1];
            text += token;
          }
        }
        const content = [{ type: "text", text: JSON.parse(`"${text}"`) }];
        if (toolCall) {
          content.push(toolCall);
        }
        yield {
          content: content,
        };
      }
    } catch (error) {
      console.error("MyModelAdapter run error:", error);
      throw error;
    }
  },
};

Does anyone know how to resolve this issue?

[rennokki]

I think it has to do with this: https://github.com/vercel/ai/issues/5493

For some reason, Assistant UI is not handling the deltas right. @Yonom?

[rennokki]

I opened https://github.com/assistant-ui/assistant-ui/issues/1901 and then closed it, it seems to be from the upstream. Gonna come back if I don't find anything relevant.

[rennokki]

I think it has to do with continueSteps: https://github.com/vercel/ai/issues/6125

[ShobhitPatra]

I've reproduced and analyzed this issue. The assistant-ui library handles accumulation correctly – the problem is in the custom adapter implementation.

The bug: Content array is recreated on each chunk, losing previous tool calls.

// Wrong - replaces content
const toolCalls = chunk.toolCalls || [];  // Only current chunk
const content = [{ type: "text", text }];
yield { content };  // Previous tool calls lost

The fix: Accumulate tool calls outside the loop.

// Correct - accumulates content
const toolCallsMap = new Map();

for (const chunk of stream) {
  for (const tc of chunk.toolCalls || []) {
    toolCallsMap.set(tc.id, tc);
  }
  
  yield {
    content: [
      { type: "text", text },
      ...Array.from(toolCallsMap.values())
    ]
  };
}

Reproduction: https://github.com/ShobhitPatra/assistant-ui-reproducible-issue-1718