diff --git a/packages/twenty-front/src/modules/ai/components/__tests__/AiChatAssistantMessageRenderer.test.tsx b/packages/twenty-front/src/modules/ai/components/__tests__/AiChatAssistantMessageRenderer.test.tsx
index a7f22ec15c..6ac31eb56e 100644
--- a/packages/twenty-front/src/modules/ai/components/__tests__/AiChatAssistantMessageRenderer.test.tsx
+++ b/packages/twenty-front/src/modules/ai/components/__tests__/AiChatAssistantMessageRenderer.test.tsx
@@ -8,12 +8,14 @@ jest.mock('@/ai/components/ThinkingStepsDisplay', () => ({
ThinkingStepsDisplay: ({
hasAssistantTextResponseStarted,
parts,
+ isTrailingWhileStreaming,
}: {
parts: unknown[];
hasAssistantTextResponseStarted: boolean;
+ isTrailingWhileStreaming?: boolean;
}) => (
- {`thinking-${parts.length}-${hasAssistantTextResponseStarted ? 'answer-started' : 'answer-pending'}`}
+ {`thinking-${parts.length}-${hasAssistantTextResponseStarted ? 'answer-started' : 'answer-pending'}${isTrailingWhileStreaming ? '-trailing-while-streaming' : ''}`}
),
}));
@@ -40,12 +42,15 @@ jest.mock('@/ai/components/CodeExecutionDisplay', () => ({
CodeExecutionDisplay: () => ,
}));
-const renderAssistantRenderer = (messageParts: ExtendedUIMessagePart[]) => {
+const renderAssistantRenderer = (
+ messageParts: ExtendedUIMessagePart[],
+ { isLastMessageStreaming = false }: { isLastMessageStreaming?: boolean } = {},
+) => {
return render(
,
);
@@ -234,6 +239,70 @@ describe('AiChatAssistantMessageRenderer', () => {
expect(screen.getByTestId('code-execution-display')).toBeInTheDocument();
});
+ it('should flag the trailing thinking steps group while streaming', () => {
+ const messageParts = [
+ {
+ type: 'reasoning',
+ text: 'Reasoning content',
+ state: 'done',
+ },
+ {
+ type: 'tool-web_search',
+ toolCallId: 'tool-1',
+ input: { query: 'crm software' },
+ output: { result: { ok: true } },
+ state: 'output-available',
+ },
+ ] as ExtendedUIMessagePart[];
+
+ renderAssistantRenderer(messageParts, { isLastMessageStreaming: true });
+
+ expect(screen.getByTestId('thinking-steps-display')).toHaveTextContent(
+ 'trailing-while-streaming',
+ );
+ });
+
+ it('should not flag a thinking steps group when answer text follows it', () => {
+ const messageParts = [
+ {
+ type: 'tool-web_search',
+ toolCallId: 'tool-1',
+ input: { query: 'crm software' },
+ output: { result: { ok: true } },
+ state: 'output-available',
+ },
+ {
+ type: 'text',
+ text: 'Partial answer',
+ state: 'streaming',
+ },
+ ] as ExtendedUIMessagePart[];
+
+ renderAssistantRenderer(messageParts, { isLastMessageStreaming: true });
+
+ expect(screen.getByTestId('thinking-steps-display')).not.toHaveTextContent(
+ 'trailing-while-streaming',
+ );
+ });
+
+ it('should not flag the trailing thinking steps group when the message is not streaming', () => {
+ const messageParts = [
+ {
+ type: 'tool-web_search',
+ toolCallId: 'tool-1',
+ input: { query: 'crm software' },
+ output: { result: { ok: true } },
+ state: 'output-available',
+ },
+ ] as ExtendedUIMessagePart[];
+
+ renderAssistantRenderer(messageParts);
+
+ expect(screen.getByTestId('thinking-steps-display')).not.toHaveTextContent(
+ 'trailing-while-streaming',
+ );
+ });
+
it('should group a dynamic-tool part (native web search) into ThinkingStepsDisplay', () => {
const messageParts = [
{
diff --git a/packages/twenty-front/src/modules/ai/components/__tests__/ThinkingStepsDisplay.test.tsx b/packages/twenty-front/src/modules/ai/components/__tests__/ThinkingStepsDisplay.test.tsx
index 2e4d68695b..47d10bfb4b 100644
--- a/packages/twenty-front/src/modules/ai/components/__tests__/ThinkingStepsDisplay.test.tsx
+++ b/packages/twenty-front/src/modules/ai/components/__tests__/ThinkingStepsDisplay.test.tsx
@@ -64,28 +64,32 @@ const createReasoningPart = ({
const createToolPart = ({
input = { query: 'crm software' },
output = { result: { ok: true } },
+ state = 'output-available',
type = 'tool-web_search',
}: {
type?: `tool-${string}`;
input?: Record;
output?: unknown;
+ state?: string;
} = {}): ThinkingStepPart =>
({
type,
toolCallId: `${type}-call-id`,
input,
output,
- state: 'output-available',
+ state,
}) as ThinkingStepPart;
const renderThinkingStepsDisplay = ({
hasAssistantTextResponseStarted = false,
isLastMessageStreaming,
parts,
+ isTrailingWhileStreaming = false,
}: {
parts: ThinkingStepPart[];
isLastMessageStreaming: boolean;
hasAssistantTextResponseStarted?: boolean;
+ isTrailingWhileStreaming?: boolean;
}) => {
return render(
@@ -93,6 +97,7 @@ const renderThinkingStepsDisplay = ({
parts={parts}
isLastMessageStreaming={isLastMessageStreaming}
hasAssistantTextResponseStarted={hasAssistantTextResponseStarted}
+ isTrailingWhileStreaming={isTrailingWhileStreaming}
/>
,
);
@@ -120,6 +125,49 @@ describe('ThinkingStepsDisplay', () => {
expect(document.querySelector('svg[viewBox="0 0 14 14"]')).not.toBeNull();
});
+ it('should render the loading label for a tool step awaiting its output while streaming', () => {
+ renderThinkingStepsDisplay({
+ isLastMessageStreaming: true,
+ parts: [createToolPart({ output: null, state: 'input-available' })],
+ });
+
+ expect(
+ screen.getByText('Searching the web for crm software'),
+ ).toBeInTheDocument();
+ });
+
+ it('should append the pending thinking row after completed steps when requested', () => {
+ renderThinkingStepsDisplay({
+ isLastMessageStreaming: true,
+ isTrailingWhileStreaming: true,
+ parts: [createToolPart()],
+ });
+
+ expect(screen.getByText('Thinking')).toBeInTheDocument();
+ });
+
+ it('should not render a thinking row for completed steps by default', () => {
+ renderThinkingStepsDisplay({
+ isLastMessageStreaming: true,
+ parts: [createToolPart()],
+ });
+
+ expect(screen.queryByText('Thinking')).toBeNull();
+ });
+
+ it('should not append the pending thinking row while a tool step is still running', () => {
+ renderThinkingStepsDisplay({
+ isLastMessageStreaming: true,
+ isTrailingWhileStreaming: true,
+ parts: [createToolPart({ output: null, state: 'input-available' })],
+ });
+
+ expect(screen.queryByText('Thinking')).toBeNull();
+ expect(
+ screen.getByText('Searching the web for crm software'),
+ ).toBeInTheDocument();
+ });
+
it('should render done state collapsed by default', () => {
renderThinkingStepsDisplay({
isLastMessageStreaming: false,
diff --git a/packages/twenty-front/src/modules/ai/utils/__tests__/thinkingStepsDisplayState.test.ts b/packages/twenty-front/src/modules/ai/utils/__tests__/thinkingStepsDisplayState.test.ts
index b094a8114b..2c0ea8a129 100644
--- a/packages/twenty-front/src/modules/ai/utils/__tests__/thinkingStepsDisplayState.test.ts
+++ b/packages/twenty-front/src/modules/ai/utils/__tests__/thinkingStepsDisplayState.test.ts
@@ -23,12 +23,14 @@ const createToolPart = ({
errorText,
input = {},
output,
+ state = 'output-available',
type = 'tool-web_search',
}: {
type?: `tool-${string}`;
input?: Record;
output?: unknown;
errorText?: string;
+ state?: string;
} = {}): ThinkingStepPart =>
({
type,
@@ -36,7 +38,7 @@ const createToolPart = ({
input,
output,
errorText,
- state: 'output-available',
+ state,
}) as ThinkingStepPart;
describe('thinkingStepsDisplayState', () => {
@@ -91,27 +93,32 @@ describe('thinkingStepsDisplayState', () => {
expect(isThinkingStepPartActive(reasoningPart, false)).toBe(true);
});
- it('should mark tool parts without output as active while message is streaming', () => {
+ it('should mark tool parts awaiting their output as active while message is streaming', () => {
const toolPart = createToolPart({
type: 'tool-web_search',
- output: undefined,
- errorText: undefined,
+ state: 'input-available',
});
expect(isThinkingStepPartActive(toolPart, true)).toBe(true);
expect(isThinkingStepPartActive(toolPart, false)).toBe(false);
});
- it('should mark tool parts with output or error as inactive', () => {
+ it('should mark completed and failed tool parts as inactive', () => {
const completedToolPart = createToolPart({
output: { result: { ok: true } },
});
+ const completedNullOutputToolPart = createToolPart({
+ output: null,
+ });
const failedToolPart = createToolPart({
- output: undefined,
errorText: 'Tool failed',
+ state: 'output-error',
});
expect(isThinkingStepPartActive(completedToolPart, true)).toBe(false);
+ expect(isThinkingStepPartActive(completedNullOutputToolPart, true)).toBe(
+ false,
+ );
expect(isThinkingStepPartActive(failedToolPart, true)).toBe(false);
});
});
diff --git a/packages/twenty-front/src/modules/ai/utils/isThinkingStepPartActive.ts b/packages/twenty-front/src/modules/ai/utils/isThinkingStepPartActive.ts
index 24c02c1957..a037724235 100644
--- a/packages/twenty-front/src/modules/ai/utils/isThinkingStepPartActive.ts
+++ b/packages/twenty-front/src/modules/ai/utils/isThinkingStepPartActive.ts
@@ -1,5 +1,3 @@
-import { isDefined } from 'twenty-shared/utils';
-
import { type ThinkingStepPart } from '@/ai/utils/thinkingStepPart';
export const isThinkingStepPartActive = (
@@ -12,7 +10,6 @@ export const isThinkingStepPartActive = (
return (
isLastMessageStreaming &&
- !isDefined(part.output) &&
- !isDefined(part.errorText)
+ (part.state === 'input-streaming' || part.state === 'input-available')
);
};