feat: support chat sessions management & recover from storage (#4399)

* feat: support multiple session management

* fix: history style

* feat: support recover chat sessions

* fix: additional data recover

* fix: use builtin icon

packages/ai-native/src/browser/ai-core.contribution.ts

@@ -73,12 +73,16 @@ import {
   AI_CHAT_VIEW_ID,
   AI_MENU_BAR_DEBUG_TOOLBAR,
   ChatProxyServiceToken,
+  IChatInternalService,
+  IChatManagerService,
   ISumiMCPServerBackend,
   SumiMCPServerProxyServicePath,
 } from '../common';
 import { MCPServerDescription } from '../common/mcp-server-manager';
 
+import { ChatManagerService } from './chat/chat-manager.service';
 import { ChatProxyService } from './chat/chat-proxy.service';
+import { ChatInternalService } from './chat/chat.internal.service';
 import { AIChatView } from './chat/chat.view';
 import { CodeActionSingleHandler } from './contrib/code-action/code-action.handler';
 import { AIInlineCompletionsProvider } from './contrib/inline-completions/completeProvider';
@@ -226,17 +230,25 @@ export class AINativeBrowserContribution
   @Autowired(WorkbenchEditorService)
   private readonly workbenchEditorService: WorkbenchEditorServiceImpl;
 
+  @Autowired(IChatManagerService)
+  private readonly chatManagerService: ChatManagerService;
+
+  @Autowired(IChatInternalService)
+  private readonly chatInternalService: ChatInternalService;
+
   constructor() {
     this.registerFeature();
   }
 
-  initialize() {
+  async initialize() {
     const { supportsChatAssistant } = this.aiNativeConfigService.capabilities;
 
     if (supportsChatAssistant) {
       ComponentRegistryImpl.addLayoutModule(this.appConfig.layoutConfig, AI_CHAT_VIEW_ID, AI_CHAT_CONTAINER_ID);
       ComponentRegistryImpl.addLayoutModule(this.appConfig.layoutConfig, DESIGN_MENU_BAR_RIGHT, AI_CHAT_LOGO_AVATAR_ID);
       this.chatProxyService.registerDefaultAgent();
+      this.chatInternalService.init();
+      await this.chatManagerService.init();
     }
   }
 

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

@@ -4,32 +4,101 @@ import {
   CancellationTokenSource,
   Disposable,
   DisposableMap,
+  Emitter,
   IChatProgress,
+  IStorage,
+  STORAGE_NAMESPACE,
+  StorageProvider,
 } from '@opensumi/ide-core-common';
-import { ChatMessageRole } from '@opensumi/ide-core-common/lib/types/ai-native';
-import { IChatMessage } from '@opensumi/ide-core-common/lib/types/ai-native';
-
-import { IChatAgentService } from '../../common';
-
-import { ChatModel, ChatRequestModel } from './chat-model';
+import { ChatMessageRole, IChatMessage, IHistoryChatMessage } from '@opensumi/ide-core-common/lib/types/ai-native';
+
+import { IChatAgentService, IChatFollowup, IChatRequestMessage, IChatResponseErrorDetails } from '../../common';
+import { MsgHistoryManager } from '../model/msg-history-manager';
+
+import { ChatModel, ChatRequestModel, ChatResponseModel, IChatProgressResponseContent } from './chat-model';
+
+interface ISessionModel {
+  sessionId: string;
+  history: { additional: Record<string, any>; messages: IHistoryChatMessage[] };
+  requests: {
+    requestId: string;
+    message: IChatRequestMessage;
+    response: {
+      isCanceled: boolean;
+      responseText: string;
+      responseContents: IChatProgressResponseContent[];
+      errorDetails: IChatResponseErrorDetails | undefined;
+      followups: IChatFollowup[];
+    };
+  }[];
+}
 
 @Injectable()
 export class ChatManagerService extends Disposable {
   #sessionModels = this.registerDispose(new DisposableMap<string, ChatModel>());
   #pendingRequests = this.registerDispose(new DisposableMap<string, CancellationTokenSource>());
+  private storageInitEmitter = new Emitter<void>();
+  public onStorageInit = this.storageInitEmitter.event;
 
   @Autowired(INJECTOR_TOKEN)
   injector: Injector;
 
   @Autowired(IChatAgentService)
   chatAgentService: IChatAgentService;
 
+  @Autowired(StorageProvider)
+  private storageProvider: StorageProvider;
+
+  private _chatStorage: IStorage;
+
+  protected fromJSON(data: ISessionModel[]) {
+    // TODO: 支持ApplyService恢复
+    return data.map((item) => {
+      const model = new ChatModel({
+        sessionId: item.sessionId,
+        history: new MsgHistoryManager(item.history),
+      });
+      const requests = item.requests.map(
+        (request) =>
+          new ChatRequestModel(
+            request.requestId,
+            model,
+            request.message,
+            new ChatResponseModel(request.requestId, model, request.message.agentId, {
+              responseContents: request.response.responseContents,
+              isComplete: true,
+              responseText: request.response.responseText,
+              errorDetails: request.response.errorDetails,
+              followups: request.response.followups,
+              isCanceled: request.response.isCanceled,
+            }),
+          ),
+      );
+      model.restoreRequests(requests);
+      return model;
+    });
+  }
+
   constructor() {
     super();
   }
 
+  async init() {
+    this._chatStorage = await this.storageProvider(STORAGE_NAMESPACE.CHAT);
+    const sessionsModelData = this._chatStorage.get<ISessionModel[]>('sessionModels', []);
+    const savedSessions = this.fromJSON(sessionsModelData);
+    savedSessions.forEach((session) => {
+      this.#sessionModels.set(session.sessionId, session);
+    });
+    await this.storageInitEmitter.fireAndAwait();
+  }
+
+  getSessions() {
+    return Array.from(this.#sessionModels.values());
+  }
+
   startSession() {
-    const model = this.injector.get(ChatModel);
+    const model = new ChatModel();
     this.#sessionModels.set(model.sessionId, model);
     return model;
   }
@@ -46,6 +115,7 @@ export class ChatManagerService extends Disposable {
     this.#sessionModels.disposeKey(sessionId);
     this.#pendingRequests.get(sessionId)?.cancel();
     this.#pendingRequests.disposeKey(sessionId);
+    this.saveSessions();
   }
 
   createRequest(sessionId: string, message: string, agentId: string, command?: string) {
@@ -122,11 +192,17 @@ export class ChatManagerService extends Disposable {
     } finally {
       listener.dispose();
       this.#pendingRequests.disposeKey(model.sessionId);
+      this.saveSessions();
     }
   }
 
+  protected saveSessions() {
+    this._chatStorage.set('sessionModels', this.getSessions());
+  }
+
   cancelRequest(sessionId: string) {
     this.#pendingRequests.get(sessionId)?.cancel();
     this.#pendingRequests.disposeKey(sessionId);
+    this.saveSessions();
   }
 }

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

@@ -80,9 +80,29 @@ export class ChatResponseModel extends Disposable {
     return this.#onDidChange.event;
   }
 
-  constructor(requestId: string, public readonly session: IChatModel, public readonly agentId: string) {
+  constructor(
+    requestId: string,
+    public readonly session: IChatModel,
+    public readonly agentId: string,
+    initParams?: {
+      isComplete: boolean;
+      isCanceled: boolean;
+      responseContents: IChatProgressResponseContent[];
+      responseText: string;
+      errorDetails: IChatResponseErrorDetails | undefined;
+      followups: IChatFollowup[] | undefined;
+    },
+  ) {
     super();
     this.#requestId = requestId;
+    if (initParams) {
+      this.#responseContents = initParams.responseContents;
+      this.#responseText = initParams.responseText;
+      this.#isComplete = initParams.isComplete;
+      this.#isCanceled = initParams.isCanceled;
+      this.#errorDetails = initParams.errorDetails;
+      this.#followups = initParams.followups;
+    }
   }
 
   updateContent(progress: IChatProgress, quiet?: boolean): void {
@@ -215,6 +235,16 @@ export class ChatResponseModel extends Disposable {
     this.#followups = followups;
     this.#onDidChange.fire();
   }
+
+  toJSON() {
+    return {
+      isCanceled: this.isCanceled,
+      responseContents: this.responseContents,
+      responseText: this.responseText,
+      errorDetails: this.errorDetails,
+      followups: this.followups,
+    };
+  }
 }
 
 export class ChatRequestModel implements IChatRequestModel {
@@ -231,19 +261,29 @@ export class ChatRequestModel implements IChatRequestModel {
   ) {
     this.#requestId = requestId;
   }
+
+  toJSON() {
+    return {
+      requestId: this.requestId,
+      message: this.message,
+      response: this.response,
+    };
+  }
 }
 
-@Injectable({ multiple: true })
 export class ChatModel extends Disposable implements IChatModel {
   private static requestIdPool = 0;
 
-  @Autowired(ILogger)
-  protected readonly logger: ILogger;
-
-  @Autowired(INJECTOR_TOKEN)
-  injector: Injector;
+  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 = uuid();
+  #sessionId: string;
   get sessionId(): string {
     return this.#sessionId;
   }
@@ -253,11 +293,12 @@ export class ChatModel extends Disposable implements IChatModel {
     return Array.from(this.#requests.values());
   }
 
-  @memoize
-  get history(): MsgHistoryManager {
-    return this.injector.get(MsgHistoryManager, []);
+  restoreRequests(requests: ChatRequestModel[]): void {
+    this.#requests = new Map(requests.map((r) => [r.requestId, r]));
   }
 
+  readonly history: MsgHistoryManager;
+
   addRequest(message: IChatRequestMessage): ChatRequestModel {
     const requestId = `${this.sessionId}_request_${ChatModel.requestIdPool++}`;
     const response = new ChatResponseModel(requestId, this, message.agentId);
@@ -279,7 +320,8 @@ export class ChatModel extends Disposable implements IChatModel {
     if (basicKind.includes(kind)) {
       request.response.updateContent(progress, quiet);
     } else {
-      this.logger.error(`Couldn't handle progress: ${JSON.stringify(progress)}`);
+      // eslint-disable-next-line no-console
+      console.error(`Couldn't handle progress: ${JSON.stringify(progress)}`);
     }
   }
 
@@ -291,6 +333,14 @@ export class ChatModel extends Disposable implements IChatModel {
     super.dispose();
     this.#requests.forEach((r) => r.response.dispose());
   }
+
+  toJSON() {
+    return {
+      sessionId: this.sessionId,
+      history: this.history,
+      requests: this.requests,
+    };
+  }
 }
 
 @Injectable({ multiple: true })

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

@@ -39,7 +39,7 @@ import { ChatFeatureRegistry } from './chat.feature.registry';
 @Injectable()
 export class ChatProxyService extends Disposable {
   // 避免和插件注册的 agent id 冲突
-  static readonly AGENT_ID = 'Default_Chat_Agent_' + uuid(6);
+  static readonly AGENT_ID = 'Default_Chat_Agent';
 
   @Autowired(IChatAgentService)
   private readonly chatAgentService: IChatAgentService;

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

@@ -43,9 +43,15 @@ export class ChatInternalService extends Disposable {
     return this.#sessionModel;
   }
 
-  constructor() {
-    super();
-    this.#sessionModel = this.chatManagerService.startSession();
+  init() {
+    this.chatManagerService.onStorageInit(() => {
+      const sessions = this.chatManagerService.getSessions();
+      if (sessions.length > 0) {
+        this.activateSession(sessions[sessions.length - 1].sessionId);
+      } else {
+        this.createSessionModel();
+      }
+    });
   }
 
   public setLatestRequestId(id: string): void {
@@ -70,12 +76,24 @@ export class ChatInternalService extends Disposable {
     this._onCancelRequest.fire();
   }
 
-  clearSessionModel() {
-    this.chatManagerService.clearSession(this.#sessionModel.sessionId);
+  createSessionModel() {
     this.#sessionModel = this.chatManagerService.startSession();
     this._onChangeSession.fire(this.#sessionModel.sessionId);
   }
 
+  clearSessionModel(sessionId?: string) {
+    sessionId = sessionId || this.#sessionModel.sessionId;
+    this.chatManagerService.clearSession(sessionId);
+    if (sessionId === this.#sessionModel.sessionId) {
+      this.#sessionModel = this.chatManagerService.startSession();
+    }
+    this._onChangeSession.fire(this.#sessionModel.sessionId);
+  }
+
+  getSessions() {
+    return this.chatManagerService.getSessions();
+  }
+
   activateSession(sessionId: string) {
     const targetSession = this.chatManagerService.getSession(sessionId);
     if (!targetSession) {