From ec65d29983c5812c3a945a438f1b914027ab471b Mon Sep 17 00:00:00 2001
From: Stefan Dirix <sdirix@eclipsesource.com>
Date: Fri, 24 Jan 2025 11:50:12 +0100
Subject: [PATCH] chore: further improvements

- streamline parameters for change set deletion
- remove the ref for latestRequest again
---
 .../src/browser/chat-input-widget.tsx         | 23 +++++++++++--------
 .../src/browser/chat-view-widget.tsx          |  8 +++----
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/packages/ai-chat-ui/src/browser/chat-input-widget.tsx b/packages/ai-chat-ui/src/browser/chat-input-widget.tsx
index dcce6ae468167..9b01001804b52 100644
--- a/packages/ai-chat-ui/src/browser/chat-input-widget.tsx
+++ b/packages/ai-chat-ui/src/browser/chat-input-widget.tsx
@@ -137,8 +137,8 @@ export class AIChatInputWidget extends ReactWidget {
 interface ChatInputProperties {
     onCancel: (requestModel: ChatRequestModel) => void;
     onQuery: (query: string) => void;
-    onDeleteChangeSet: (requestModel: ChatRequestModel) => void;
-    onDeleteChangeSetElement: (requestModel: ChatRequestModel, index: number) => void;
+    onDeleteChangeSet: (sessionId: string) => void;
+    onDeleteChangeSetElement: (sessionId: string, index: number) => void;
     isEnabled?: boolean;
     chatModel: ChatModel;
     editorProvider: MonacoEditorProvider;
@@ -150,10 +150,14 @@ interface ChatInputProperties {
 }
 
 const ChatInput: React.FunctionComponent<ChatInputProperties> = (props: ChatInputProperties) => {
+    const onDeleteChangeSet = () => props.onDeleteChangeSet(props.chatModel.id);
+    const onDeleteChangeSetElement = (index: number) => props.onDeleteChangeSetElement(props.chatModel.id, index);
+
     const [inProgress, setInProgress] = React.useState(false);
     const [changeSetUI, setChangeSetUI] = React.useState(
         () => props.chatModel.changeSet ? buildChangeSetUI(props.chatModel.changeSet, props.labelProvider, onDeleteChangeSet, onDeleteChangeSetElement) : undefined
     );
+
     // eslint-disable-next-line no-null/no-null
     const editorContainerRef = React.useRef<HTMLDivElement | null>(null);
     // eslint-disable-next-line no-null/no-null
@@ -235,7 +239,6 @@ const ChatInput: React.FunctionComponent<ChatInputProperties> = (props: ChatInpu
         };
     }, []);
 
-    const latestRequest = React.useRef<ChatRequestModel>();
     const responseListenerRef = React.useRef<Disposable>();
     // track chat model updates to keep our UI in sync
     // - keep "inProgress" in sync with the request state
@@ -243,7 +246,6 @@ const ChatInput: React.FunctionComponent<ChatInputProperties> = (props: ChatInpu
     React.useEffect(() => {
         const listener = props.chatModel.onDidChange(event => {
             if (event.kind === 'addRequest') {
-                latestRequest.current = event.request;
                 if (event.request) {
                     setInProgress(ChatRequestModel.isInProgress(event.request));
                 }
@@ -259,6 +261,7 @@ const ChatInput: React.FunctionComponent<ChatInputProperties> = (props: ChatInpu
                 }
             }
         });
+        setChangeSetUI(props.chatModel.changeSet ? buildChangeSetUI(props.chatModel.changeSet, props.labelProvider, onDeleteChangeSet, onDeleteChangeSetElement) : undefined);
         return () => {
             listener?.dispose();
             responseListenerRef.current?.dispose();
@@ -316,15 +319,13 @@ const ChatInput: React.FunctionComponent<ChatInputProperties> = (props: ChatInpu
         className: 'codicon-add'
     }] : [];
 
-    const onDeleteChangeSet = () => latestRequest.current ? props.onDeleteChangeSet(latestRequest.current) : undefined;
-    const onDeleteChangeSetElement = (index: number) => latestRequest.current ? props.onDeleteChangeSetElement(latestRequest.current, index) : undefined;
-
     const rightOptions = inProgress
         ? [{
             title: 'Cancel (Esc)',
             handler: () => {
-                if (latestRequest.current) {
-                    props.onCancel(latestRequest.current);
+                const latestRequest = getLatestRequest(props.chatModel);
+                if (latestRequest) {
+                    props.onCancel(latestRequest);
                 }
                 setInProgress(false);
             },
@@ -484,3 +485,7 @@ function acceptablePendingElements(changeSet: ChangeSet): ChangeSetElement[] {
     return changeSet.getElements().filter(e => e.accept && (e.state === undefined || e.state === 'pending'));
 }
 
+function getLatestRequest(chatModel: ChatModel): ChatRequestModel | undefined {
+    const requests = chatModel.getRequests();
+    return requests.length > 0 ? requests[requests.length - 1] : undefined;
+}
diff --git a/packages/ai-chat-ui/src/browser/chat-view-widget.tsx b/packages/ai-chat-ui/src/browser/chat-view-widget.tsx
index 66e6920dbcc37..c368971297460 100644
--- a/packages/ai-chat-ui/src/browser/chat-view-widget.tsx
+++ b/packages/ai-chat-ui/src/browser/chat-view-widget.tsx
@@ -181,12 +181,12 @@ export class ChatViewWidget extends BaseWidget implements ExtractableWidget, Sta
         this.chatService.cancelRequest(requestModel.session.id, requestModel.id);
     }
 
-    protected onDeleteChangeSet(requestModel: ChatRequestModel): void {
-        this.chatService.deleteChangeSet(requestModel.session.id);
+    protected onDeleteChangeSet(sessionId: string): void {
+        this.chatService.deleteChangeSet(sessionId);
     }
 
-    protected onDeleteChangeSetElement(requestModel: ChatRequestModel, index: number): void {
-        this.chatService.deleteChangeSetElement(requestModel.session.id, index);
+    protected onDeleteChangeSetElement(sessionId: string, index: number): void {
+        this.chatService.deleteChangeSetElement(sessionId, index);
     }
 
     lock(): void {