chore: improve chat context (#4407)

* chore: improve chat context

* Update packages/ai-native/src/browser/components/ChatInput.tsx

Co-authored-by: LouisLv <zheyanglv@qq.com>

* chore: update session manager

---------

Co-authored-by: LouisLv <zheyanglv@qq.com>

Author: Aaaaash

packages/ai-native/__test__/browser/chat/chat-agent.service.test.ts

@@ -1,9 +1,12 @@
-import { CancellationToken } from '@opensumi/ide-core-common';
+import { CancellationToken, Emitter } from '@opensumi/ide-core-common';
+import { ChatFeatureRegistryToken, ChatServiceToken } from '@opensumi/ide-core-common/lib/types/ai-native';
 import { createBrowserInjector } from '@opensumi/ide-dev-tool/src/injector-helper';
 import { MockInjector } from '@opensumi/ide-dev-tool/src/mock-injector';
 
 import { ChatAgentService } from '../../../lib/browser/chat/chat-agent.service';
 import { IChatAgent, IChatAgentMetadata, IChatAgentRequest, IChatManagerService } from '../../../lib/common';
+import { LLMContextServiceToken } from '../../../lib/common/llm-context';
+import { ChatAgentPromptProvider } from '../../../lib/common/prompts/context-prompt-provider';
 
 describe('ChatAgentService', () => {
   let injector: MockInjector;
@@ -19,9 +22,32 @@ describe('ChatAgentService', () => {
             startSession: jest.fn(),
           },
         },
+        {
+          token: ChatAgentPromptProvider,
+          useValue: {
+            provideContextPrompt: (val, msg) => msg,
+          },
+        },
+        {
+          token: ChatServiceToken,
+          useValue: {
+          },
+        },
+        {
+          token: LLMContextServiceToken,
+          useValue: {
+            onDidContextFilesChangeEvent: new Emitter().event,
+            serialize: () => { },
+          },
+        },
+        {
+          token: ChatFeatureRegistryToken,
+          useValue: {
+          },
+        },
       ]),
     );
-    chatAgentService = new ChatAgentService();
+    chatAgentService = injector.get(ChatAgentService);
   });
 
   it('should register an agent', () => {
@@ -42,7 +68,7 @@ describe('ChatAgentService', () => {
       id: 'agent1',
       metadata: {},
       provideSlashCommands: () => Promise.resolve([]),
-      invoke: () => {},
+      invoke: () => { },
     } as unknown as IChatAgent;
     chatAgentService.registerAgent(agent);
 

packages/ai-native/src/browser/chat/chat-agent.service.ts

@@ -4,6 +4,7 @@ import { Autowired, Injectable } from '@opensumi/di';
 import {
   CancellationToken,
   ChatFeatureRegistryToken,
+  ChatMessageRole,
   ChatServiceToken,
   Disposable,
   Emitter,
@@ -24,6 +25,8 @@ import {
   IChatFollowup,
   IChatMessageStructure,
 } from '../../common';
+import { LLMContextService, LLMContextServiceToken } from '../../common/llm-context';
+import { ChatAgentPromptProvider } from '../../common/prompts/context-prompt-provider';
 import { IChatFeatureRegistry } from '../types';
 
 import { ChatService } from './chat.api.service';
@@ -34,6 +37,12 @@ export class ChatAgentService extends Disposable implements IChatAgentService {
 
   private defaultAgentId: string | undefined;
 
+  private initialUserMessageMap: Map<string, string> = new Map();
+
+  private shouldUpdateContext = false;
+
+  private contextVersion: number;
+
   private readonly _onDidChangeAgents = new Emitter<void>();
   readonly onDidChangeAgents = this._onDidChangeAgents.event;
 
@@ -43,6 +52,12 @@ export class ChatAgentService extends Disposable implements IChatAgentService {
   @Autowired(ILogger)
   logger: ILogger;
 
+  @Autowired(LLMContextServiceToken)
+  protected readonly contextService: LLMContextService;
+
+  @Autowired(ChatAgentPromptProvider)
+  protected readonly promptProvider: ChatAgentPromptProvider;
+
   @Autowired(ChatServiceToken)
   private aiChatService: ChatService;
 
@@ -52,6 +67,12 @@ export class ChatAgentService extends Disposable implements IChatAgentService {
   constructor() {
     super();
     this.addDispose(this._onDidChangeAgents);
+    this.addDispose(this.contextService.onDidContextFilesChangeEvent((event) => {
+      if (event.version !== this.contextVersion) {
+        this.contextVersion = event.version;
+        this.shouldUpdateContext = true;
+      }
+    }));
   }
 
   registerAgent(agent: IChatAgent): IDisposable {
@@ -119,10 +140,28 @@ export class ChatAgentService extends Disposable implements IChatAgentService {
       throw new Error(`No agent with id ${id}`);
     }
 
+    // 发送第一条消息时携带初始 context
+    if (!this.initialUserMessageMap.has(request.sessionId)) {
+      this.initialUserMessageMap.set(request.sessionId, request.message);
+      const rawMessage = request.message;
+      request.message = this.provideContextMessage(rawMessage);
+    }
+
+    if (this.shouldUpdateContext) {
+      request.message = this.provideContextMessage(request.message);
+      this.shouldUpdateContext = false;
+    }
+
     const result = await data.agent.invoke(request, progress, history, token);
     return result;
   }
 
+  private provideContextMessage(message: string) {
+    const context = this.contextService.serialize();
+    const fullMessage = this.promptProvider.provideContextPrompt(context, message);
+    return fullMessage;
+  }
+
   populateChatInput(id: string, message: IChatMessageStructure) {
     this.aiChatService.sendMessage({
       ...message,

packages/ai-native/src/browser/chat/chat-model.ts

@@ -1,4 +1,4 @@
-import { Autowired, INJECTOR_TOKEN, Injectable, Injector } from '@opensumi/di';
+import { Injectable } from '@opensumi/di';
 import {
   Disposable,
   Emitter,
@@ -8,8 +8,6 @@ import {
   IChatProgress,
   IChatToolContent,
   IChatTreeData,
-  ILogger,
-  memoize,
   uuid,
 } from '@opensumi/ide-core-common';
 import { MarkdownString, isMarkdownString } from '@opensumi/monaco-editor-core/esm/vs/base/common/htmlContent';
@@ -274,13 +272,18 @@ export class ChatRequestModel implements IChatRequestModel {
 export class ChatModel extends Disposable implements IChatModel {
   private static requestIdPool = 0;
 
-  constructor(initParams?: { sessionId?: string; history?: MsgHistoryManager; requests?: ChatRequestModel[] }) {
+  constructor(initParams?: {
+    sessionId?: string;
+    history?: MsgHistoryManager;
+    requests?: ChatRequestModel[];
+  }) {
     super();
     this.#sessionId = initParams?.sessionId ?? uuid();
     this.history = initParams?.history ?? new MsgHistoryManager();
     if (initParams?.requests) {
       this.#requests = new Map(initParams.requests.map((r) => [r.requestId, r]));
     }
+
   }
 
   #sessionId: string;
@@ -300,9 +303,11 @@ export class ChatModel extends Disposable implements IChatModel {
   readonly history: MsgHistoryManager;
 
   addRequest(message: IChatRequestMessage): ChatRequestModel {
+    const msg = message;
+
     const requestId = `${this.sessionId}_request_${ChatModel.requestIdPool++}`;
-    const response = new ChatResponseModel(requestId, this, message.agentId);
-    const request = new ChatRequestModel(requestId, this, message, response);
+    const response = new ChatResponseModel(requestId, this, msg.agentId);
+    const request = new ChatRequestModel(requestId, this, msg, response);
 
     this.#requests.set(requestId, request);
     return request;

packages/ai-native/src/browser/chat/chat-proxy.service.ts

@@ -12,7 +12,7 @@ import {
   IAIReporter,
   IApplicationService,
   IChatProgress,
-  uuid,
+  getOperatingSystemName,
 } from '@opensumi/ide-core-common';
 import { AINativeSettingSectionsId } from '@opensumi/ide-core-common/lib/settings/ai-native';
 import { IChatMessage } from '@opensumi/ide-core-common/lib/types/ai-native';
@@ -92,7 +92,7 @@ export class ChatProxyService extends Disposable {
               AINativeSettingSectionsId.SystemPrompt,
               'You are a powerful AI coding assistant working in OpenSumi, a top IDE framework. You collaborate with a USER to solve coding tasks, which may involve creating, modifying, or debugging code, or answering questions. When the USER sends a message, relevant context (e.g., open files, cursor position, edit history, linter errors) may be attached. Use this information as needed.\n\n<tool_calling>\nYou have access to tools to assist with tasks. Follow these rules:\n1. Always adhere to the tool call schema and provide all required parameters.\n2. Only use tools explicitly provided; ignore unavailable ones.\n3. Avoid mentioning tool names to the USER (e.g., say "I will edit your file" instead of "I need to use the edit_file tool").\n4. Only call tools when necessary; respond directly if the task is general or you already know the answer.\n5. Explain to the USER why you’re using a tool before calling it.\n</tool_calling>\n\n<making_code_changes>\nWhen modifying code:\n1. Use code edit tools instead of outputting code unless explicitly requested.\n2. Limit tool calls to one per turn.\n3. Ensure generated code is immediately executable by including necessary imports, dependencies, and endpoints.\n4. For new projects, create a dependency management file (e.g., requirements.txt) and a README.\n5. For web apps, design a modern, user-friendly UI.\n6. Avoid generating non-textual or excessively long code.\n7. Read file contents before editing, unless appending a small change or creating a new file.\n8. Fix introduced linter errors if possible, but stop after 3 attempts and ask the USER for guidance.\n9. Reapply reasonable code edits if they weren’t followed initially.\n</making_code_changes>\n\nUse the appropriate tools to fulfill the USER’s request, ensuring all required parameters are provided or inferred from context.',
             ) +
-            `\n\n<user_info>\nThe user's OS version is ${this.applicationService.frontendOS}. The absolute path of the user's workspace is ${this.appConfig.workspaceDir}.\n</user_info>`,
+            `\n\n<user_info>\nThe user's OS is ${getOperatingSystemName()}. The absolute path of the user's workspace is ${this.appConfig.workspaceDir}.\n</user_info>`,
         },
         invoke: async (
           request: IChatAgentRequest,

packages/ai-native/src/browser/chat/chat.module.less

@@ -279,7 +279,6 @@
   }
 }
 
-
 .chat_tips_container {
   display: flex;
   align-items: center;
@@ -290,6 +289,6 @@
 }
 
 .chat_history {
-  width: calc(100% - 60px);
+  width: calc(100% - 40px);
   color: var(--design-text-foreground);
 }

packages/ai-native/src/browser/chat/chat.view.tsx

@@ -5,7 +5,6 @@ import {
   AINativeConfigService,
   CommandService,
   getIcon,
-  useEventEffect,
   useInjectable,
   useUpdateOnEvent,
 } from '@opensumi/ide-core-browser';
@@ -42,8 +41,6 @@ import {
   IChatMessageStructure,
   TokenMCPServerProxyService,
 } from '../../common';
-import { LLMContextService, LLMContextServiceToken } from '../../common/llm-context';
-import { ChatAgentPromptProvider } from '../../common/prompts/context-prompt-provider';
 import { ChatContext } from '../components/ChatContext';
 import { CodeBlockWrapperInput } from '../components/ChatEditor';
 import ChatHistory, { IChatHistoryItem } from '../components/ChatHistory';
@@ -82,18 +79,12 @@ export const AIChatView = () => {
   const chatAgentService = useInjectable<IChatAgentService>(IChatAgentService);
   const chatFeatureRegistry = useInjectable<ChatFeatureRegistry>(ChatFeatureRegistryToken);
   const chatRenderRegistry = useInjectable<ChatRenderRegistry>(ChatRenderRegistryToken);
-  const contextService = useInjectable<LLMContextService>(LLMContextServiceToken);
-  const promptProvider = useInjectable<ChatAgentPromptProvider>(ChatAgentPromptProvider);
-  const mcpServerProxyService = useInjectable<MCPServerProxyService>(TokenMCPServerProxyService);
 
   const layoutService = useInjectable<IMainLayoutService>(IMainLayoutService);
   const msgHistoryManager = aiChatService.sessionModel.history;
   const containerRef = React.useRef<HTMLDivElement>(null);
   const autoScroll = React.useRef<boolean>(true);
   const chatInputRef = React.useRef<{ setInputValue: (v: string) => void } | null>(null);
-  const dialogService = useInjectable<IDialogService>(IDialogService);
-  const aiNativeConfigService = useInjectable<AINativeConfigService>(AINativeConfigService);
-  const commandService = useInjectable<CommandService>(CommandService);
 
   const [shortcutCommands, setShortcutCommands] = React.useState<ChatSlashCommandItemModel[]>([]);
 
@@ -115,8 +106,6 @@ export const AIChatView = () => {
   const [defaultAgentId, setDefaultAgentId] = React.useState<string>('');
   const [command, setCommand] = React.useState('');
   const [theme, setTheme] = React.useState<string | null>(null);
-  const [mcpToolsCount, setMcpToolsCount] = React.useState<number>(0);
-  const [mcpServersCount, setMcpServersCount] = React.useState<number>(0);
 
   React.useEffect(() => {
     const featureSlashCommands = chatFeatureRegistry.getAllShortcutSlashCommand();
@@ -537,10 +526,7 @@ export const AIChatView = () => {
       const { message, agentId, command, reportExtra } = value;
       const { actionType, actionSource } = reportExtra || {};
 
-      const context = contextService.serialize();
-      const fullMessage = await promptProvider.provideContextPrompt(context, message);
-
-      const request = aiChatService.createRequest(fullMessage, agentId!, command);
+      const request = aiChatService.createRequest(message, agentId!, command);
       if (!request) {
         return;
       }
@@ -689,32 +675,6 @@ export const AIChatView = () => {
     };
   }, [aiChatService.sessionModel]);
 
-  useEventEffect(
-    mcpServerProxyService.onChangeMCPServers,
-    () => {
-      mcpServerProxyService.getAllMCPTools().then((tools) => {
-        setMcpToolsCount(tools.length);
-      });
-      mcpServerProxyService.$getServers().then((servers) => {
-        setMcpServersCount(servers.length);
-      });
-    },
-    [mcpServerProxyService],
-  );
-
-  const handleShowMCPTools = React.useCallback(async () => {
-    const tools = await mcpServerProxyService.getAllMCPTools();
-    dialogService.open({
-      message: <MCPToolsDialog tools={tools} />,
-      type: MessageType.Empty,
-      buttons: ['关闭'],
-    });
-  }, [mcpServerProxyService, dialogService]);
-
-  const handleShowMCPConfig = React.useCallback(() => {
-    commandService.executeCommand(OPEN_MCP_CONFIG_COMMAND.id);
-  }, [commandService]);
-
   return (
     <div id={styles.ai_chat_view}>
       <div className={styles.header_container}>
@@ -754,18 +714,6 @@ export const AIChatView = () => {
                   </Popover>
                 ))}
               </div>
-              <div className={styles.header_operate_right}>
-                {aiNativeConfigService.capabilities.supportsMCP && (
-                  <>
-                    <div className={styles.tag} onClick={handleShowMCPConfig}>
-                      {`MCP Servers: ${mcpServersCount}`}
-                    </div>
-                    <div className={styles.tag} onClick={handleShowMCPTools}>
-                      {`MCP Tools: ${mcpToolsCount}`}
-                    </div>
-                  </>
-                )}
-              </div>
             </div>
             <ChatInputWrapperRender
               onSend={(value, agentId, command) =>
@@ -808,6 +756,8 @@ export function DefaultChatViewHeader({
   const aiNativeConfigService = useInjectable<AINativeConfigService>(AINativeConfigService);
   const mcpServerProxyService = useInjectable<MCPServerProxyService>(TokenMCPServerProxyService);
   const aiChatService = useInjectable<ChatInternalService>(IChatInternalService);
+  const commandService = useInjectable<CommandService>(CommandService);
+
   const [historyList, setHistoryList] = React.useState<IChatHistoryItem[]>([]);
   const [currentTitle, setCurrentTitle] = React.useState<string>('');
   const handleNewChat = React.useCallback(() => {
@@ -828,6 +778,10 @@ export function DefaultChatViewHeader({
     [aiChatService],
   );
 
+  const handleShowMCPConfig = React.useCallback(() => {
+    commandService.executeCommand(OPEN_MCP_CONFIG_COMMAND.id);
+  }, [commandService]);
+
   const handleShowMCPTools = React.useCallback(async () => {
     const tools = await mcpServerProxyService.getAllMCPTools();
     dialogService.open({
@@ -916,23 +870,6 @@ export function DefaultChatViewHeader({
           ariaLabel={localize('aiNative.operate.clear.title')}
         />
       </Popover>
-      {aiNativeConfigService.capabilities.supportsMCP && (
-        <Popover
-          overlayClassName={styles.popover_icon}
-          id={'ai-chat-header-tools'}
-          position={PopoverPosition.left}
-          title={localize('aiNative.operate.tools.title')}
-        >
-          <EnhanceIcon
-            wrapperClassName={styles.action_btn}
-            className={getIcon('menubar-tool')}
-            onClick={handleShowMCPTools}
-            tabIndex={0}
-            role='button'
-            ariaLabel={localize('aiNative.operate.tools.title')}
-          />
-        </Popover>
-      )}
       <Popover
         overlayClassName={styles.popover_icon}
         id={'ai-chat-header-close'}

packages/ai-native/src/browser/components/ChatContext/index.tsx

@@ -39,7 +39,7 @@ export const ChatContext = memo(() => {
       50,
     )((files) => {
       if (files) {
-        updateAddedFiles(files);
+        updateAddedFiles([...files.attached]);
       }
     }, contextService);
 
@@ -57,7 +57,7 @@ export const ChatContext = memo(() => {
   }, []);
 
   const onDidDeselect = useCallback((uri: URI) => {
-    contextService.removeFileFromContext(uri);
+    contextService.removeFileFromContext(uri, true);
   }, []);
 
   const onDidClickFile = useCallback((uri: URI) => {