feat: add message compaction for AI chats (#19205)

2026-04-02 20:43:11 +05:30
parent c1e4756f9c
commit 7f2b853ae1
11 changed files with 146 additions and 15 deletions
@@ -1,3 +1,4 @@
+import { AIChatCompactionIndicator } from '@/ai/components/AIChatCompactionIndicator';
 import { CodeExecutionDisplay } from '@/ai/components/CodeExecutionDisplay';
 import { RoutingStatusDisplay } from '@/ai/components/RoutingStatusDisplay';
 import { ThinkingStepsDisplay } from '@/ai/components/ThinkingStepsDisplay';
@@ -58,6 +59,8 @@ const MessagePartRenderer = ({
      return <LazyMarkdownRenderer text={part.text} />;
    case 'data-routing-status':
      return <RoutingStatusDisplay data={part.data} />;
+    case 'data-compaction':
+      return <AIChatCompactionIndicator />;
    case 'data-code-execution':
      return (
        <CodeExecutionDisplay
@@ -0,0 +1,30 @@
+import { styled } from '@linaria/react';
+import { useContext } from 'react';
+import { IconTransform } from 'twenty-ui/display';
+import { ThemeContext, themeCssVariables } from 'twenty-ui/theme-constants';
+
+const StyledIndicatorContainer = styled.div`
+  align-items: center;
+  color: ${themeCssVariables.font.color.tertiary};
+  display: flex;
+  gap: ${themeCssVariables.spacing[1]};
+`;
+
+const StyledIconTextContainer = styled.div`
+  align-items: center;
+  display: flex;
+  gap: ${themeCssVariables.spacing[1]};
+`;
+
+export const AIChatCompactionIndicator = () => {
+  const { theme } = useContext(ThemeContext);
+
+  return (
+    <StyledIndicatorContainer>
+      <StyledIconTextContainer>
+        <IconTransform size={theme.icon.size.sm} />
+        <div>The conversation has been compacted</div>
+      </StyledIconTextContainer>
+    </StyledIndicatorContainer>
+  );
+};
@@ -51,7 +51,7 @@ export const mapDBPartToUIMessagePart = (
      };
    case 'step-start':
      return {
-        type: 'step-start',
+        type: part.type,
      };
    case 'data-routing-status':
      return {
@@ -63,6 +63,8 @@ export const mapUIMessagePartsToDBParts = (
          };
        case 'step-start':
          return basePart;
+        case 'data-compaction':
+          return null;
        case 'data-routing-status':
          return {
            ...basePart,
@@ -44,6 +44,7 @@ import { AgentChatStreamingService } from './services/agent-chat-streaming.servi
 import { AgentChatService } from './services/agent-chat.service';
 import { AgentTitleGenerationService } from './services/agent-title-generation.service';
 import { ChatExecutionService } from './services/chat-execution.service';
+import { MessagePruningService } from './services/message-pruning.service';
 import { SystemPromptBuilderService } from './services/system-prompt-builder.service';

@Module({
@@ -110,6 +111,7 @@ import { SystemPromptBuilderService } from './services/system-prompt-builder.ser
    AgentChatStreamingService,
    AgentTitleGenerationService,
    ChatExecutionService,
+    MessagePruningService,
    StreamAgentChatJob,
    SystemPromptBuilderService,
  ],
@@ -17,4 +17,5 @@ export type StreamAgentChatJobData = {
  lastUserMessageParts: ExtendedUIMessagePart[];
  hasTitle: boolean;
  existingTurnId?: string;
+  conversationSizeTokens: number;
 };
@@ -9,23 +9,23 @@ import type {
 } from 'twenty-shared/ai';
 import { Repository } from 'typeorm';

-import { AgentChatCancelSubscriberService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-cancel-subscriber.service';
-import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.service';
-import { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util';
-import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
 import { Process } from 'src/engine/core-modules/message-queue/decorators/process.decorator';
 import { Processor } from 'src/engine/core-modules/message-queue/decorators/processor.decorator';
 import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants';
+import { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util';
+import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
 import { AgentMessageRole } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity';
-import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service';
 import { computeCostBreakdown } from 'src/engine/metadata-modules/ai/ai-billing/utils/compute-cost-breakdown.util';
 import { convertDollarsToBillingCredits } from 'src/engine/metadata-modules/ai/ai-billing/utils/convert-dollars-to-billing-credits.util';
 import { extractCacheCreationTokens } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util';
-import type { AIModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type';
 import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
+import { AgentChatCancelSubscriberService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-cancel-subscriber.service';
+import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.service';
+import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service';
 import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service';
 import { ChatExecutionService } from 'src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service';
 import { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util';
+import type { AIModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type';

 import { STREAM_AGENT_CHAT_JOB_NAME } from './stream-agent-chat-job-name.constant';
 import { type StreamAgentChatJobData } from './stream-agent-chat-job.types';
@@ -212,6 +212,14 @@ export class StreamAgentChatJob {
            });
          };

+          const onCompaction = () => {
+            writer.write({
+              type: 'data-compaction' as const,
+              id: `compaction-${data.threadId}`,
+              data: {},
+            });
+          };
+
          const { stream, modelConfig } =
            await this.chatExecutionService.streamChat({
              workspace,
@@ -220,7 +228,9 @@ export class StreamAgentChatJob {
              browsingContext: data.browsingContext,
              modelId: data.modelId,
              onCodeExecutionUpdate,
+              onCompaction,
              abortSignal,
+              conversationSizeTokens: data.conversationSizeTokens,
            });

          const titleWritePromise = titlePromise.then((generatedTitle) => {
@@ -8,16 +8,16 @@ import { InjectMessageQueue } from 'src/engine/core-modules/message-queue/decora
 import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants';
 import { MessageQueueService } from 'src/engine/core-modules/message-queue/services/message-queue.service';
 import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
-import {
-  AgentException,
-  AgentExceptionCode,
-} from 'src/engine/metadata-modules/ai/ai-agent/agent.exception';
-import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
 import {
  AgentMessageRole,
  AgentMessageStatus,
 } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity';
 import { mapDBPartsToUIMessageParts } from 'src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartsToUIMessageParts';
+import {
+  AgentException,
+  AgentExceptionCode,
+} from 'src/engine/metadata-modules/ai/ai-agent/agent.exception';
+import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
 import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
 import { STREAM_AGENT_CHAT_JOB_NAME } from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job-name.constant';
 import { type StreamAgentChatJobData } from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job.types';
@@ -99,6 +99,7 @@ export class AgentChatStreamingService {
        lastUserMessageText: text,
        lastUserMessageParts: [{ type: 'text', text }],
        hasTitle: !!thread.title,
+        conversationSizeTokens: thread.conversationSize,
        existingTurnId: savedUserMessage.turnId ?? undefined,
      },
    );
@@ -155,7 +156,10 @@ export class AgentChatStreamingService {
      event: { type: 'message-persisted', messageId: nextQueued.id },
    });

-    const uiMessages = await this.loadMessagesFromDB(threadId, userWorkspaceId);
+    const [uiMessages, thread] = await Promise.all([
+      this.loadMessagesFromDB(threadId, userWorkspaceId),
+      this.threadRepository.findOneByOrFail({ id: threadId }),
+    ]);

    const streamId = generateId();

@@ -171,6 +175,7 @@ export class AgentChatStreamingService {
        lastUserMessageText: messageText,
        lastUserMessageParts: [{ type: 'text', text: messageText }],
        hasTitle,
+        conversationSizeTokens: thread.conversationSize,
        existingTurnId: turnId,
      },
    );
@@ -40,6 +40,7 @@ import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agen
 import { repairToolCall } from 'src/engine/metadata-modules/ai/ai-agent/utils/repair-tool-call.util';
 import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service';
 import { extractCacheCreationTokensFromSteps } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util';
+import { MessagePruningService } from 'src/engine/metadata-modules/ai/ai-chat/services/message-pruning.service';
 import { SystemPromptBuilderService } from 'src/engine/metadata-modules/ai/ai-chat/services/system-prompt-builder.service';
 import {
  extractCodeInterpreterFiles,
@@ -51,12 +52,12 @@ import {
  AI_SDK_OPENAI,
 } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-sdk-package.const';
 import { AI_TELEMETRY_CONFIG } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-telemetry.const';
-import { type AIModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type';
 import {
  AiModelRegistryService,
  type RegisteredAIModel,
 } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service';
 import { SdkProviderFactoryService } from 'src/engine/metadata-modules/ai/ai-models/services/sdk-provider-factory.service';
+import { type AIModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type';
 import { SkillService } from 'src/engine/metadata-modules/skill/skill.service';

 export type ChatExecutionOptions = {
@@ -65,8 +66,10 @@ export type ChatExecutionOptions = {
  messages: UIMessage<unknown, UIDataTypes, UITools>[];
  browsingContext: BrowsingContextType | null;
  onCodeExecutionUpdate?: CodeExecutionStreamEmitter;
+  onCompaction?: () => void;
  modelId?: string;
  abortSignal?: AbortSignal;
+  conversationSizeTokens: number;
 };

 export type ChatExecutionResult = {
@@ -89,6 +92,7 @@ export class ChatExecutionService {
    private readonly systemPromptBuilder: SystemPromptBuilderService,
    private readonly exceptionHandlerService: ExceptionHandlerService,
    private readonly sdkProviderFactory: SdkProviderFactoryService,
+    private readonly messagePruningService: MessagePruningService,
  ) {}

  async streamChat({
@@ -97,8 +101,10 @@ export class ChatExecutionService {
    messages,
    browsingContext,
    onCodeExecutionUpdate,
+    onCompaction,
    modelId,
    abortSignal,
+    conversationSizeTokens,
  }: ChatExecutionOptions): Promise<ChatExecutionResult> {
    const { actorContext, roleId, userId, userContext } =
      await this.agentActorContextService.buildUserAndAgentActorContext(
@@ -232,7 +238,26 @@ export class ChatExecutionService {
            : undefined,
    };

-    const modelMessages = await convertToModelMessages(processedMessages);
+    const rawModelMessages = await convertToModelMessages(processedMessages);
+
+    const pruningResult =
+      this.messagePruningService.pruneIfOverContextWindowLimit(
+        rawModelMessages,
+        modelConfig.contextWindowTokens,
+        conversationSizeTokens,
+      );
+
+    if (pruningResult.isStillOverLimit) {
+      throw new Error(
+        'This conversation is too long for the model to process. Please start a new thread.',
+      );
+    }
+
+    if (pruningResult.wasPruned) {
+      onCompaction?.();
+    }
+
+    const modelMessages = pruningResult.messages;

    const billUsageFromSteps = (steps: StepResult<ToolSet>[]) => {
      const usage = steps.reduce<LanguageModelUsage>(
@@ -0,0 +1,52 @@
+import { Injectable, Logger } from '@nestjs/common';
+
+import { type ModelMessage, pruneMessages } from 'ai';
+
+const COMPACTION_THRESHOLD_RATIO = 0.9;
+const TOOL_CALLS_PRESERVE_LAST_N_MESSAGES = 2;
+
+export type PruningResult = {
+  messages: ModelMessage[];
+  wasPruned: boolean;
+  isStillOverLimit: boolean;
+};
+
+@Injectable()
+export class MessagePruningService {
+  private readonly logger = new Logger(MessagePruningService.name);
+
+  pruneIfOverContextWindowLimit(
+    messages: ModelMessage[],
+    contextWindowTokens: number,
+    conversationSizeTokens: number,
+  ): PruningResult {
+    const threshold = contextWindowTokens * COMPACTION_THRESHOLD_RATIO;
+
+    if (conversationSizeTokens < threshold) {
+      return { messages, wasPruned: false, isStillOverLimit: false };
+    }
+
+    this.logger.log(
+      `Conversation size ${conversationSizeTokens} exceeds threshold ${Math.round(threshold)} (${contextWindowTokens} * ${COMPACTION_THRESHOLD_RATIO}). Pruning messages.`,
+    );
+
+    const prunedMessages = pruneMessages({
+      messages,
+      reasoning: 'before-last-message',
+      toolCalls: `before-last-${TOOL_CALLS_PRESERVE_LAST_N_MESSAGES}-messages`,
+      emptyMessages: 'remove',
+    });
+
+    const wasPruned = prunedMessages.length < messages.length;
+
+    if (wasPruned) {
+      this.logger.log(
+        `Pruned ${messages.length - prunedMessages.length} messages (${messages.length} → ${prunedMessages.length})`,
+      );
+    }
+
+    const isStillOverLimit = !wasPruned;
+
+    return { messages: prunedMessages, wasPruned, isStillOverLimit };
+  }
+}
@@ -88,4 +88,5 @@ export type DataMessagePart = {
  };
  'code-execution': CodeExecutionData;
  'thread-title': { title: string };
+  compaction: Record<string, never>;
 };