edit ability

src/vs/workbench/contrib/void/browser/chatThreadService.ts

@@ -40,6 +40,7 @@ export type ChatMessage =
 		content: string | null; // content sent to the llm - allowed to be '', will be replaced with (empty)
 		displayContent: string | null; // content displayed to user  - allowed to be '', will be ignored
 		selections: StagingSelectionItem[] | null; // the user's selection
+		stagingSelections: StagingSelectionItem[] | null; // staging selections in edit mode
 	}
 	| {
 		role: 'assistant';
@@ -59,13 +60,14 @@ export type ChatThreads = {
 		createdAt: string; // ISO string
 		lastModified: string; // ISO string
 		messages: ChatMessage[];
+		stagingSelections: StagingSelectionItem[] | null;
+		focusedMessageIdx?: number | undefined; // index of the message that is being edited (undefined if none)
 	};
 }
 
 export type ThreadsState = {
 	allThreads: ChatThreads;
 	currentThreadId: string; // intended for internal use only
-	currentStagingSelections: StagingSelectionItem[] | null;
 }
 
 export type ThreadStreamState = {
@@ -84,6 +86,8 @@ const newThreadObject = () => {
 		createdAt: now,
 		lastModified: now,
 		messages: [],
+		focusedMessageIdx: undefined,
+		stagingSelections: null,
 	} satisfies ChatThreads[string]
 }
 
@@ -105,8 +109,12 @@ export interface IChatThreadService {
 	openNewThread(): void;
 	switchToThread(threadId: string): void;
 
-	setStaging(stagingSelection: StagingSelectionItem[] | null): void;
+	getFocusedMessageIdx(): number | undefined;
+	setFocusedMessageIdx(messageIdx: number | undefined): void;
 
+	_useStagingSelectionsState(messageIdx?: number | undefined): readonly [StagingSelectionItem[], (selections: StagingSelectionItem[]) => void];
+
+	editUserMessageAndStreamResponse(userMessage: string, messageIdx: number): Promise<void>;
 	addUserMessageAndStreamResponse(userMessage: string): Promise<void>;
 	cancelStreaming(threadId: string): void;
 	dismissStreamError(threadId: string): void;
@@ -137,7 +145,6 @@ class ChatThreadService extends Disposable implements IChatThreadService {
 		this.state = {
 			allThreads: this._readAllThreads(),
 			currentThreadId: null as unknown as string, // gets set in startNewThread()
-			currentStagingSelections: null,
 		}
 
 		// always be in a thread
@@ -187,18 +194,50 @@ class ChatThreadService extends Disposable implements IChatThreadService {
 		this._setStreamState(threadId, { messageSoFar: undefined, streamingToken: undefined, error })
 	}
 
-	async addUserMessageAndStreamResponse(userMessage: string) {
-		const threadId = this.getCurrentThread().id
 
-		const currSelns = this.state.currentStagingSelections ?? []
+	async editUserMessageAndStreamResponse(userMessage: string, messageIdx: number) {
+
+		const thread = this.getCurrentThread()
+
+		const messageToReplace = thread.messages[messageIdx]
+		if (messageToReplace?.role !== 'user') {
+			console.log(`Error: tried to edit non-user message. messageIdx=${messageIdx}, numMessages=${thread.messages.length}`)
+			return
+		}
+
+		// clear messages up to the index
+		const slicedMessages = thread.messages.slice(0, messageIdx)
+		this._setState({
+			allThreads: {
+				...this.state.allThreads,
+				[thread.id]: {
+					...thread,
+					messages: slicedMessages
+				}
+			}
+		}, true)
+
+		// stream the edit
+		this.addUserMessageAndStreamResponse(userMessage, messageToReplace.stagingSelections)
+
+	}
+
+	async addUserMessageAndStreamResponse(userMessage: string, selectionsOverride?: StagingSelectionItem[] | null) {
+
+
+		const thread = this.getCurrentThread()
+		const threadId = thread.id
+
+		let defaultThreadSelections = thread.stagingSelections
+
+		const currSelns = selectionsOverride ?? defaultThreadSelections ?? [] // don't use _useFocusedStagingState to avoid race conditions with focusing
 
 		// add user's message to chat history
 		const instructions = userMessage
 		const content = await chat_userMessage(instructions, currSelns, this._modelService)
-		const userHistoryElt: ChatMessage = { role: 'user', content: content, displayContent: instructions, selections: currSelns }
+		const userHistoryElt: ChatMessage = { role: 'user', content: content, displayContent: instructions, selections: currSelns, stagingSelections: [], }
 		this._addMessageToThread(threadId, userHistoryElt)
 
-
 		this._setStreamState(threadId, { error: undefined })
 
 		const llmCancelToken = this._llmMessageService.sendLLMMessage({
@@ -210,12 +249,15 @@ class ChatThreadService extends Disposable implements IChatThreadService {
 				...this.getCurrentThread().messages.map(m => ({ role: m.role, content: m.content || '(null)' })),
 			],
 			onText: ({ newText, fullText }) => {
+				console.log('onText', fullText)
 				this._setStreamState(threadId, { messageSoFar: fullText })
 			},
 			onFinalMessage: ({ fullText: content }) => {
+				console.log('finalMessage', JSON.stringify(content))
 				this.finishStreaming(threadId, content)
 			},
 			onError: (error) => {
+				console.log('onError', content)
 				this.finishStreaming(threadId, this.streamState[threadId]?.messageSoFar ?? '', error)
 			},
 
@@ -241,7 +283,12 @@ class ChatThreadService extends Disposable implements IChatThreadService {
 
 	getCurrentThread(): ChatThreads[string] {
 		const state = this.state
-		return state.allThreads[state.currentThreadId];
+		return state.allThreads[state.currentThreadId]
+	}
+
+	getFocusedMessageIdx() {
+		const thread = this.getCurrentThread()
+		return thread.focusedMessageIdx
 	}
 
 	switchToThread(threadId: string) {
@@ -291,11 +338,93 @@ class ChatThreadService extends Disposable implements IChatThreadService {
 		this._setState({ allThreads: newThreads }, true) // the current thread just changed (it had a message added to it)
 	}
 
+	// sets the currently selected message (must be undefined if no message is selected)
+	setFocusedMessageIdx(messageIdx: number | undefined) {
 
-	setStaging(stagingSelection: StagingSelectionItem[] | null): void {
-		this._setState({ currentStagingSelections: stagingSelection }, true) // this is a hack for now
+		const threadId = this.state.currentThreadId
+		const thread = this.state.allThreads[threadId]
+		if (!thread) return
+
+		this._setState({
+			allThreads: {
+				...this.state.allThreads,
+				[threadId]: {
+					...thread,
+					focusedMessageIdx: messageIdx
+				}
+			}
+		}, true)
 	}
 
+	// set thread.messages[messageIdx].stagingSelections
+	private setEditMessageStagingSelections(stagingSelections: StagingSelectionItem[], messageIdx: number): void {
+
+		const thread = this.getCurrentThread()
+		const message = thread.messages[messageIdx]
+		if (message.role !== 'user') return;
+
+		this._setState({
+			allThreads: {
+				...this.state.allThreads,
+				[thread.id]: {
+					...thread,
+					messages: thread.messages.map((m, i) =>
+						i === messageIdx ? { ...m, stagingSelections } : m
+					)
+				}
+			}
+		}, true)
+
+	}
+
+	// set thread.stagingSelections
+	private setDefaultStagingSelections(stagingSelections: StagingSelectionItem[]): void {
+
+
+		console.log('Default1')
+		const thread = this.getCurrentThread()
+
+		console.log('Default2')
+
+		this._setState({
+			allThreads: {
+				...this.state.allThreads,
+				[thread.id]: {
+					...thread,
+					stagingSelections
+				}
+			}
+		}, true)
+
+	}
+
+	// gets `staging` and `setStaging` of the currently focused element, given the index of the currently selected message (or undefined if no message is selected)
+	_useStagingSelectionsState(messageIdx?: number | undefined) {
+
+		let staging: StagingSelectionItem[] = []
+		let setStaging: (selections: StagingSelectionItem[]) => void = () => { }
+
+		const thread = this.getCurrentThread()
+		const isFocusingMessage = messageIdx !== undefined
+		if (isFocusingMessage) { // is editing message
+
+			const message = thread.messages[messageIdx!]
+			if (message.role === 'user') {
+				staging = message.stagingSelections || []
+				setStaging = (s) => this.setEditMessageStagingSelections(s, messageIdx)
+			}
+
+		}
+		else { // is editing the default input box
+			staging = thread.stagingSelections || []
+			setStaging = this.setDefaultStagingSelections.bind(this)
+		}
+
+		return [staging, setStaging] as const
+	}
+
+
+
 }
 
 registerSingleton(IChatThreadService, ChatThreadService, InstantiationType.Eager);

src/vs/workbench/contrib/void/browser/react/src/sidebar-tsx/SidebarChat.tsx

@@ -25,6 +25,7 @@ import { Pencil } from 'lucide-react';
 import { FeatureName } from '../../../../../../../platform/void/common/voidSettingsTypes.js';
 
 
+
 export const IconX = ({ size, className = '', ...props }: { size: number, className?: string } & React.SVGProps<SVGSVGElement>) => {
 	return (
 		<svg
@@ -152,6 +153,7 @@ interface VoidInputFormProps {
 	showModelDropdown?: boolean;
 	showSelections?: boolean;
 	selections?: any[];
+	showProspectiveSelections?: boolean;
 	onSelectionsChange?: (selections: any[]) => void;
 
 	// Optional close button
@@ -171,6 +173,7 @@ export const VoidInputForm: React.FC<VoidInputFormProps> = ({
 	featureName,
 	showSelections = false,
 	selections = [],
+	showProspectiveSelections = true,
 	onSelectionsChange,
 }) => {
 	return (
@@ -191,6 +194,7 @@ export const VoidInputForm: React.FC<VoidInputFormProps> = ({
 					type='staging'
 					selections={selections}
 					setSelections={onSelectionsChange}
+					showProspectiveSelections={showProspectiveSelections}
 				/>
 			)}
 
@@ -528,14 +532,61 @@ export const SelectedFiles = (
 
 
 type ChatBubbleMode = 'display' | 'edit'
-const ChatBubble = ({ chatMessage, isLoading }: { chatMessage: ChatMessage, isLoading?: boolean, }) => {
+const ChatBubble = ({ chatMessage, isLoading, messageIdx }: { chatMessage: ChatMessage, messageIdx?: number, isLoading?: boolean, }) => {
 
 	const role = chatMessage.role
 
+
+	const accessor = useAccessor()
+	const chatThreadsService = accessor.get('IChatThreadService')
+
 	// edit mode state
+	const [staging, setStaging] = chatThreadsService._useStagingSelectionsState(messageIdx)
 	const [mode, setMode] = useState<ChatBubbleMode>('display')
-	const [editText, setEditText] = useState(chatMessage.displayContent ?? '')
 	const [isHovered, setIsHovered] = useState(false)
+	const textAreaRef = useRef<HTMLTextAreaElement | null>(null)
+	const textAreaFnsRef = useRef<TextAreaFns | null>(null)
+	const [isDisabled, setIsDisabled] = useState(false)
+	useEffect(() => {
+		if (role === 'user') {
+			setStaging(chatMessage.selections || [])
+			if (textAreaFnsRef.current)
+				textAreaFnsRef.current.setValue(chatMessage.displayContent || '')
+			if (textAreaRef.current)
+				textAreaRef.current.value = chatMessage.displayContent || ''
+		}
+	}, [role])
+
+
+	const onSubmit = async () => {
+
+		if (isDisabled) return;
+		if (!textAreaRef.current) return;
+		if (messageIdx === undefined) return;
+
+		// reset visual state
+		setMode('display');
+
+		// stream the edit
+		const userMessage = textAreaRef.current.value;
+		await chatThreadsService.editUserMessageAndStreamResponse(userMessage, messageIdx)
+
+		textAreaFnsRef.current?.setValue('');
+	}
+
+	const onAbort = () => {
+		const threadId = chatThreadsService.state.currentThreadId
+		chatThreadsService.cancelStreaming(threadId)
+	}
+
+	const onKeyDown = useCallback((e: KeyboardEvent<HTMLTextAreaElement>) => {
+		if (e.key === 'Escape') {
+			setMode('display')
+		}
+		if (e.key === 'Enter' && !e.shiftKey) {
+			onSubmit()
+		}
+	}, [onSubmit, setMode])
 
 	if (!chatMessage.content && !isLoading) { // don't show if empty and not loading (if loading, want to show)
 		return null
@@ -552,17 +603,27 @@ const ChatBubble = ({ chatMessage, isLoading }: { chatMessage: ChatMessage, isLo
 		}
 		else if (mode === 'edit') {
 			chatbubbleContents = <>
-				<SelectedFiles type='past' selections={chatMessage.selections || []} />
-				<textarea
-					value={editText}
-					onChange={(e) => setEditText(e.target.value)}
-					className={`
-						w-full max-w-full
-						h-auto min-h-[81px] max-h-[500px]
-						bg-void-bg-1 resize-none
-					`}
-					style={{ marginTop: 0 }}
-				/>
+				<VoidInputForm
+					onSubmit={onSubmit}
+					onAbort={onAbort}
+					isStreaming={false}
+					isDisabled={isDisabled}
+					showSelections={true}
+					selections={staging || []}
+					showProspectiveSelections={false}
+					onSelectionsChange={setStaging}
+					featureName="Ctrl+L"
+				>
+					<VoidInputBox2
+						className='min-h-[81px] max-h-[500px] p-1'
+						placeholder="Edit your message..."
+						onChangeText={(text) => setIsDisabled(!text)}
+						onKeyDown={onKeyDown}
+						ref={textAreaRef}
+						fnsRef={textAreaFnsRef}
+						multiline={true}
+					/>
+				</VoidInputForm>
 			</>
 		}
 	}
@@ -606,7 +667,10 @@ const ChatBubble = ({ chatMessage, isLoading }: { chatMessage: ChatMessage, isLo
 				transition-opacity duration-200 ease-in-out
 				${isHovered ? 'opacity-100' : 'opacity-0'}
 			`}
-			onClick={() => setMode(m => m === 'display' ? 'edit' : 'display')}
+			onClick={() => {
+				setMode(m => m === 'display' ? 'edit' : 'display')
+				chatThreadsService.setFocusedMessageIdx(messageIdx)
+			}}
 		/>}
 	</div>
 }
@@ -641,7 +705,7 @@ export const SidebarChat = () => {
 
 	const currentThread = chatThreadsService.getCurrentThread()
 	const previousMessages = currentThread?.messages ?? []
-	const selections = chatThreadsState.currentStagingSelections
+	const [selections, setSelections] = chatThreadsService._useStagingSelectionsState()
 
 	// stream state
 	const currThreadStreamState = useChatThreadsStreamState(chatThreadsState.currentThreadId)
@@ -679,7 +743,7 @@ export const SidebarChat = () => {
 		await chatThreadsService.addUserMessageAndStreamResponse(userMessage)
 		console.log('done streaming',)
 
-		chatThreadsService.setStaging([]) // clear staging
+		setSelections([]) // clear staging
 		console.log('set staging',)
 		textAreaFnsRef.current?.setValue('')
 		console.log('set value',)
@@ -687,7 +751,7 @@ export const SidebarChat = () => {
 		console.log('textAreaRef', textAreaRef.current)
 		console.log('focus',)
 
-	}, [chatThreadsService, isDisabled, isStreaming, textAreaRef, textAreaFnsRef])
+	}, [chatThreadsService, isDisabled, isStreaming, textAreaRef, textAreaFnsRef, setSelections])
 
 	const onAbort = () => {
 		const threadId = currentThread.id
@@ -708,7 +772,7 @@ export const SidebarChat = () => {
 
 	const prevMessagesHTML = useMemo(() => {
 		return previousMessages.map((message, i) =>
-			<ChatBubble key={i} chatMessage={message} />
+			<ChatBubble key={i} chatMessage={message} messageIdx={i} />
 		)
 	}, [previousMessages])
 
@@ -773,7 +837,9 @@ export const SidebarChat = () => {
 			isDisabled={isDisabled}
 			showSelections={true}
 			selections={selections || []}
-			onSelectionsChange={chatThreadsService.setStaging.bind(chatThreadsService)}
+			showProspectiveSelections={prevMessagesHTML.length === 0}
+			onSelectionsChange={setSelections}
+			// onSelectionsChange={chatThreadsService.setStagingSelections.bind(chatThreadsService)}
 			featureName="Ctrl+L"
 		>
 			<VoidInputBox2

src/vs/workbench/contrib/void/browser/react/src/util/services.tsx

@@ -288,6 +288,17 @@ export const useChatThreadsState = () => {
 		return () => { chatThreadsStateListeners.delete(ss) }
 	}, [ss])
 	return s
+	// allow user to set state natively in react
+	// const ss: React.Dispatch<React.SetStateAction<ThreadsState>> = (action)=>{
+	// 	_ss(action)
+	// 	if (typeof action === 'function') {
+	// 		const newState = action(chatThreadsState)
+	// 		chatThreadsState = newState
+	// 	} else {
+	// 		chatThreadsState = action
+	// 	}
+	// }
+	// return [s, ss] as const
 }
 
 

src/vs/workbench/contrib/void/browser/sidebarActions.ts

@@ -67,6 +67,13 @@ const getContentInRange = (model: ITextModel, range: IRange | null) => {
 }
 
 
+const findMatchingStagingIndex = (currentSelections: StagingSelectionItem[] | undefined, newSelection: StagingSelectionItem) => {
+	return currentSelections?.findIndex(s =>
+		s.fileURI.fsPath === newSelection.fileURI.fsPath
+		&& s.range?.startLineNumber === newSelection.range?.startLineNumber
+		&& s.range?.endLineNumber === newSelection.range?.endLineNumber
+	)
+}
 
 const VOID_OPEN_SIDEBAR_ACTION_ID = 'void.sidebar.open'
 registerAction2(class extends Action2 {
@@ -124,26 +131,25 @@ registerAction2(class extends Action2 {
 			range: selectionRange,
 		}
 
-		// add selection to staging
+		// update the staging selections
 		const chatThreadService = accessor.get(IChatThreadService)
-		const currentStaging = chatThreadService.state.currentStagingSelections
-		const currentStagingEltIdx = currentStaging?.findIndex(s =>
-			s.fileURI.fsPath === model.uri.fsPath
-			&& s.range?.startLineNumber === selection.range?.startLineNumber
-			&& s.range?.endLineNumber === selection.range?.endLineNumber
-		)
 
-		// if matches with existing selection, overwrite
+		const messageIdx = chatThreadService.getFocusedMessageIdx()
+		const [staging, setStaging] = chatThreadService._useStagingSelectionsState(messageIdx)
+
+		// if matches with existing selection, overwrite (since text may change)
+		const currentStagingEltIdx = findMatchingStagingIndex(staging, selection)
+
 		if (currentStagingEltIdx !== undefined && currentStagingEltIdx !== -1) {
-			chatThreadService.setStaging([
-				...currentStaging!.slice(0, currentStagingEltIdx),
+			setStaging([
+				...staging!.slice(0, currentStagingEltIdx),
 				selection,
-				...currentStaging!.slice(currentStagingEltIdx + 1, Infinity)
+				...staging!.slice(currentStagingEltIdx + 1, Infinity)
 			])
 		}
-		// if no match, add
+		// if no match, add it
 		else {
-			chatThreadService.setStaging([...(currentStaging ?? []), selection])
+			setStaging([...(staging ?? []), selection])
 		}
 
 	}