From 223943550c56f8643544ddbf8370ca458cf30085 Mon Sep 17 00:00:00 2001 From: nitin <142569587+ehconitin@users.noreply.github.com> Date: Thu, 2 Apr 2026 16:28:14 +0530 Subject: [PATCH] [AI] Unify code-interpreter streaming rendering and fix assistant width jitter (#19235) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit closes https://discord.com/channels/1130383047699738754/1480991390782455838 - Use data-code-execution as the streaming source of truth and hide duplicate code-interpreter tool parts (including tool-execute_tool wrappers). - Ensure wrapped execute_tool code-interpreter outputs still render correctly after refetch. - Gate code-interpreter server behavior by enablement state and keep assistant messages full-width to avoid streaming vs completed width shifts. Co-authored-by: Félix Malfait --- .../AIChatAssistantMessageRenderer.tsx | 10 ++-- .../modules/ai/components/AIChatMessage.tsx | 2 +- .../ai/components/ToolStepRenderer.tsx | 22 +++---- .../AIChatAssistantMessageRenderer.test.tsx | 60 ++++++++++++++++++- .../ai/utils/isCodeInterpreterToolPart.ts | 22 +++++++ .../modules/ai/utils/isThinkingStepPart.ts | 3 +- .../code-interpreter.service.ts | 10 ++++ .../providers/action-tool.provider.ts | 7 ++- .../services/chat-execution.service.ts | 21 ++++--- 9 files changed, 126 insertions(+), 31 deletions(-) create mode 100644 packages/twenty-front/src/modules/ai/utils/isCodeInterpreterToolPart.ts diff --git a/packages/twenty-front/src/modules/ai/components/AIChatAssistantMessageRenderer.tsx b/packages/twenty-front/src/modules/ai/components/AIChatAssistantMessageRenderer.tsx index ebd44920cb..9cb21641b8 100644 --- a/packages/twenty-front/src/modules/ai/components/AIChatAssistantMessageRenderer.tsx +++ b/packages/twenty-front/src/modules/ai/components/AIChatAssistantMessageRenderer.tsx @@ -6,6 +6,7 @@ import { IconDotsVertical } from 'twenty-ui/display'; import { LazyMarkdownRenderer } from '@/ai/components/LazyMarkdownRenderer'; import { ToolStepRenderer } from '@/ai/components/ToolStepRenderer'; import { groupContiguousThinkingStepParts } from '@/ai/utils/groupContiguousThinkingStepParts'; +import { isCodeInterpreterToolPart } from '@/ai/utils/isCodeInterpreterToolPart'; import { styled } from '@linaria/react'; import { isToolUIPart, type ToolUIPart } from 'ai'; import { type ExtendedUIMessagePart } from 'twenty-shared/ai'; @@ -92,16 +93,13 @@ export const AIChatAssistantMessageRenderer = ({ isLastMessageStreaming: boolean; hasError?: boolean; }) => { - // Filter out data-code-execution parts when tool-code_interpreter exists - // (the tool part contains the final result, data-code-execution is for streaming updates) - // Also filter out data-thread-title (consumed by useAgentChat, not rendered) - const hasCodeInterpreterTool = messageParts.some( - (part) => part.type === 'tool-code_interpreter', + const hasCodeExecutionData = messageParts.some( + (part) => part.type === 'data-code-execution', ); const filteredParts = messageParts.filter( (part) => part.type !== 'data-thread-title' && - (!hasCodeInterpreterTool || part.type !== 'data-code-execution'), + !(hasCodeExecutionData && isCodeInterpreterToolPart(part)), ); const renderItems = groupContiguousThinkingStepParts(filteredParts); diff --git a/packages/twenty-front/src/modules/ai/components/AIChatMessage.tsx b/packages/twenty-front/src/modules/ai/components/AIChatMessage.tsx index 564fefcd74..f023b53484 100644 --- a/packages/twenty-front/src/modules/ai/components/AIChatMessage.tsx +++ b/packages/twenty-front/src/modules/ai/components/AIChatMessage.tsx @@ -45,7 +45,7 @@ const StyledMessageText = styled.div<{ isUser?: boolean }>` padding: ${({ isUser }) => isUser ? `0 ${themeCssVariables.spacing[2]}` : '0'}; white-space: normal; - width: fit-content; + width: ${({ isUser }) => (isUser ? 'fit-content' : '100%')}; /* Pre-wrap within the whole container turns every newline between block elements into extra spacing; keep normal flow and only pre-wrap code. */ word-wrap: break-word; diff --git a/packages/twenty-front/src/modules/ai/components/ToolStepRenderer.tsx b/packages/twenty-front/src/modules/ai/components/ToolStepRenderer.tsx index 2b5378d5cb..de024008ba 100644 --- a/packages/twenty-front/src/modules/ai/components/ToolStepRenderer.tsx +++ b/packages/twenty-front/src/modules/ai/components/ToolStepRenderer.tsx @@ -8,13 +8,13 @@ import { ThemeContext, themeCssVariables } from 'twenty-ui/theme-constants'; import { CodeExecutionDisplay } from '@/ai/components/CodeExecutionDisplay'; import { ShimmeringText } from '@/ai/components/ShimmeringText'; -import { getToolIcon } from '@/ai/utils/getToolIcon'; +import { ToolOutputMessageSchema } from '@/ai/schemas/toolOutputMessageSchema'; +import { ToolOutputResultSchema } from '@/ai/schemas/toolOutputResultSchema'; import { getToolDisplayMessage, resolveToolInput, } from '@/ai/utils/getToolDisplayMessage'; -import { ToolOutputMessageSchema } from '@/ai/schemas/toolOutputMessageSchema'; -import { ToolOutputResultSchema } from '@/ai/schemas/toolOutputResultSchema'; +import { getToolIcon } from '@/ai/utils/getToolIcon'; import { useLingui } from '@lingui/react/macro'; import { type ToolUIPart } from 'ai'; import { isDefined } from 'twenty-shared/utils'; @@ -152,9 +152,15 @@ export const ToolStepRenderer = ({ const isExpandable = isDefined(output) || hasError; const ToolIcon = getToolIcon(toolName); + const outputResult = ToolOutputResultSchema.safeParse(output); + const unwrappedOutput = + rawToolName === 'execute_tool' && outputResult.success + ? outputResult.data.result + : output; + if (toolName === 'code_interpreter') { const codeInput = toolInput as { code?: string } | undefined; - const codeOutput = output as { + const codeOutput = unwrappedOutput as { result?: { stdout?: string; stderr?: string; @@ -168,7 +174,7 @@ export const ToolStepRenderer = ({ }; } | null; - const isRunning = !output && !hasError && isStreaming; + const isRunning = !unwrappedOutput && !hasError && isStreaming; return ( { ); }); - it('should keep code interpreter rendering path unchanged and out of thinking grouping', () => { + it('should show data-code-execution during streaming and hide the tool part to avoid duplicates', () => { const messageParts = [ { type: 'tool-code_interpreter', @@ -134,11 +134,65 @@ describe('AIChatAssistantMessageRenderer', () => { renderAssistantRenderer(messageParts); + expect(screen.queryByTestId('thinking-steps-display')).toBeNull(); + expect(screen.queryByTestId('tool-step-renderer')).toBeNull(); + expect(screen.getByTestId('code-execution-display')).toBeInTheDocument(); + }); + + it('should render tool-execute_tool wrapping code_interpreter via ToolStepRenderer after refetch', () => { + const messageParts = [ + { + type: 'tool-execute_tool', + toolCallId: 'tool-exec-1', + input: { + toolName: 'code_interpreter', + arguments: { code: 'print(42)' }, + }, + output: { + result: { stdout: '42', stderr: '', exitCode: 0, files: [] }, + }, + state: 'output-available', + }, + ] as ExtendedUIMessagePart[]; + + renderAssistantRenderer(messageParts); + expect(screen.queryByTestId('thinking-steps-display')).toBeNull(); expect(screen.getByTestId('tool-step-renderer')).toHaveTextContent( - 'tool-code_interpreter', + 'tool-execute_tool', ); - expect(screen.queryByTestId('code-execution-display')).toBeNull(); + }); + + it('should hide execute_tool wrapping code_interpreter when data-code-execution parts exist', () => { + const messageParts = [ + { + type: 'tool-execute_tool', + toolCallId: 'tool-exec-1', + input: { + toolName: 'code_interpreter', + arguments: { code: 'print(42)' }, + }, + output: null, + state: 'call', + }, + { + type: 'data-code-execution', + data: { + executionId: 'exec-2', + state: 'running', + code: 'print(42)', + language: 'python', + stdout: '42', + stderr: '', + files: [], + }, + }, + ] as ExtendedUIMessagePart[]; + + renderAssistantRenderer(messageParts); + + expect(screen.queryByTestId('tool-step-renderer')).toBeNull(); + expect(screen.getByTestId('code-execution-display')).toBeInTheDocument(); }); it('should render non-thinking parts directly when there are no thinking steps', () => { diff --git a/packages/twenty-front/src/modules/ai/utils/isCodeInterpreterToolPart.ts b/packages/twenty-front/src/modules/ai/utils/isCodeInterpreterToolPart.ts new file mode 100644 index 0000000000..1345993dd2 --- /dev/null +++ b/packages/twenty-front/src/modules/ai/utils/isCodeInterpreterToolPart.ts @@ -0,0 +1,22 @@ +import { isToolUIPart } from 'ai'; +import { type ExtendedUIMessagePart } from 'twenty-shared/ai'; + +export const isCodeInterpreterToolPart = ( + part: ExtendedUIMessagePart, +): boolean => { + if (!isToolUIPart(part)) { + return false; + } + + if (part.type === 'tool-code_interpreter') { + return true; + } + + if (part.type === 'tool-execute_tool') { + const input = part.input as Record | null | undefined; + + return input?.toolName === 'code_interpreter'; + } + + return false; +}; diff --git a/packages/twenty-front/src/modules/ai/utils/isThinkingStepPart.ts b/packages/twenty-front/src/modules/ai/utils/isThinkingStepPart.ts index 676246abc3..399f684edb 100644 --- a/packages/twenty-front/src/modules/ai/utils/isThinkingStepPart.ts +++ b/packages/twenty-front/src/modules/ai/utils/isThinkingStepPart.ts @@ -1,6 +1,7 @@ import { isToolUIPart } from 'ai'; import { type ExtendedUIMessagePart } from 'twenty-shared/ai'; +import { isCodeInterpreterToolPart } from '@/ai/utils/isCodeInterpreterToolPart'; import { type ThinkingStepPart } from '@/ai/utils/thinkingStepPart'; export const isThinkingStepPart = ( @@ -10,5 +11,5 @@ export const isThinkingStepPart = ( return true; } - return isToolUIPart(part) && part.type !== 'tool-code_interpreter'; + return isToolUIPart(part) && !isCodeInterpreterToolPart(part); }; diff --git a/packages/twenty-server/src/engine/core-modules/code-interpreter/code-interpreter.service.ts b/packages/twenty-server/src/engine/core-modules/code-interpreter/code-interpreter.service.ts index 4f328b2b48..a34422a8c2 100644 --- a/packages/twenty-server/src/engine/core-modules/code-interpreter/code-interpreter.service.ts +++ b/packages/twenty-server/src/engine/core-modules/code-interpreter/code-interpreter.service.ts @@ -1,6 +1,7 @@ import { Injectable } from '@nestjs/common'; import { CodeInterpreterDriverFactory } from 'src/engine/core-modules/code-interpreter/code-interpreter-driver.factory'; +import { CodeInterpreterDriverType } from 'src/engine/core-modules/code-interpreter/code-interpreter.interface'; import { type CodeExecutionResult, type CodeInterpreterDriver, @@ -8,13 +9,22 @@ import { type InputFile, type StreamCallbacks, } from 'src/engine/core-modules/code-interpreter/drivers/interfaces/code-interpreter-driver.interface'; +import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service'; @Injectable() export class CodeInterpreterService implements CodeInterpreterDriver { constructor( private readonly codeInterpreterDriverFactory: CodeInterpreterDriverFactory, + private readonly twentyConfigService: TwentyConfigService, ) {} + isEnabled(): boolean { + return ( + this.twentyConfigService.get('CODE_INTERPRETER_TYPE') !== + CodeInterpreterDriverType.DISABLED + ); + } + execute( code: string, files?: InputFile[], diff --git a/packages/twenty-server/src/engine/core-modules/tool-provider/providers/action-tool.provider.ts b/packages/twenty-server/src/engine/core-modules/tool-provider/providers/action-tool.provider.ts index 25fe033583..95428ee97b 100644 --- a/packages/twenty-server/src/engine/core-modules/tool-provider/providers/action-tool.provider.ts +++ b/packages/twenty-server/src/engine/core-modules/tool-provider/providers/action-tool.provider.ts @@ -18,6 +18,7 @@ import { type ToolDescriptor, type ToolIndexEntry, } from 'src/engine/core-modules/tool-provider/types/tool-descriptor.type'; +import { CodeInterpreterService } from 'src/engine/core-modules/code-interpreter/code-interpreter.service'; import { CodeInterpreterTool } from 'src/engine/core-modules/tool/tools/code-interpreter-tool/code-interpreter-tool'; import { DraftEmailTool } from 'src/engine/core-modules/tool/tools/email-tool/draft-email-tool'; import { SendEmailTool } from 'src/engine/core-modules/tool/tools/email-tool/send-email-tool'; @@ -41,6 +42,7 @@ export class ActionToolProvider implements ToolProvider { private readonly searchHelpCenterTool: SearchHelpCenterTool, private readonly codeInterpreterTool: CodeInterpreterTool, private readonly navigateAppTool: NavigateAppTool, + private readonly codeInterpreterService: CodeInterpreterService, private readonly permissionsService: PermissionsService, private readonly toolExecutorService: ToolExecutorService, ) { @@ -128,11 +130,12 @@ export class ActionToolProvider implements ToolProvider { ); const hasCodeInterpreterPermission = - await this.permissionsService.hasToolPermission( + this.codeInterpreterService.isEnabled() && + (await this.permissionsService.hasToolPermission( context.rolePermissionConfig, context.workspaceId, PermissionFlagType.CODE_INTERPRETER_TOOL, - ); + )); if (hasCodeInterpreterPermission) { descriptors.push( diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts index 54519d2da8..ad43d97587 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts @@ -19,6 +19,7 @@ import { UsageOperationType } from 'src/engine/core-modules/usage/enums/usage-op import { type CodeExecutionStreamEmitter } from 'src/engine/core-modules/tool-provider/interfaces/tool-provider.interface'; +import { CodeInterpreterService } from 'src/engine/core-modules/code-interpreter/code-interpreter.service'; import { WorkspaceDomainsService } from 'src/engine/core-modules/domain/workspace-domains/services/workspace-domains.service'; import { ExceptionHandlerService } from 'src/engine/core-modules/exception-handler/exception-handler.service'; import { COMMON_PRELOAD_TOOLS } from 'src/engine/core-modules/tool-provider/constants/common-preload-tools.const'; @@ -84,6 +85,7 @@ export class ChatExecutionService { private readonly aiBillingService: AiBillingService, private readonly agentActorContextService: AgentActorContextService, private readonly workspaceDomainsService: WorkspaceDomainsService, + private readonly codeInterpreterService: CodeInterpreterService, private readonly systemPromptBuilder: SystemPromptBuilderService, private readonly exceptionHandlerService: ExceptionHandlerService, private readonly sdkProviderFactory: SdkProviderFactoryService, @@ -185,19 +187,24 @@ export class ChatExecutionService { ), }; - const { processedMessages, extractedFiles } = - extractCodeInterpreterFiles(messages); + let processedMessages: UIMessage[] = messages; let storedFiles: Array<{ filename: string; fileId: string; }> = []; - if (extractedFiles.length > 0) { - storedFiles = await this.storeExtractedFiles( - extractedFiles, - workspace.id, - ); + if (this.codeInterpreterService.isEnabled()) { + const extracted = extractCodeInterpreterFiles(messages); + + processedMessages = extracted.processedMessages; + + if (extracted.extractedFiles.length > 0) { + storedFiles = await this.storeExtractedFiles( + extracted.extractedFiles, + workspace.id, + ); + } } const systemPrompt = this.systemPromptBuilder.buildFullPrompt(