Skip to content

Commit ac0e2ac

Browse files
justschenjustschen
authored
adds images in history, create file for pasted images in chat (#241664)
* references for pasted images and adding images back to history * cleanup * make sure to resize after re-reading the file * change name to specify vscode chat images * adds to workspaceStorage * delete temp storage after a whilegit add . * some cleanup
1 parent f0df737 commit ac0e2ac

File tree

4 files changed

+150
-29
lines changed

4 files changed

+150
-29
lines changed

‎src/vs/workbench/contrib/chat/browser/chatInputPart.ts

Lines changed: 51 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ import { revealInSideBarCommand } from '../../files/browser/fileActions.contribu
8080
import { ChatAgentLocation, IChatAgentService } from '../common/chatAgents.js';
8181
import { ChatContextKeys } from '../common/chatContextKeys.js';
8282
import { IChatEditingSession, WorkingSetEntryRemovalReason, WorkingSetEntryState } from '../common/chatEditingService.js';
83-
import { IChatRequestVariableEntry, isLinkVariableEntry, isPasteVariableEntry } from '../common/chatModel.js';
83+
import { IChatRequestVariableEntry, isImageVariableEntry, isLinkVariableEntry, isPasteVariableEntry } from '../common/chatModel.js';
8484
import { IChatFollowup } from '../common/chatService.js';
8585
import { IChatVariablesService } from '../common/chatVariables.js';
8686
import { IChatResponseViewModel } from '../common/chatViewModel.js';
@@ -101,6 +101,7 @@ import { ChatFollowups } from './chatFollowups.js';
101101
import { IChatViewState } from './chatWidget.js';
102102
import { ChatFileReference } from './contrib/chatDynamicVariables/chatFileReference.js';
103103
import { ChatImplicitContext } from './contrib/chatImplicitContext.js';
104+
import { resizeImage } from './imageUtils.js';
104105

105106
const $ = dom.$;
106107

@@ -538,12 +539,13 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
538539
return this.container;
539540
}
540541

541-
showPreviousValue(): void {
542+
async showPreviousValue(): Promise<void> {
542543
const inputState = this.getInputState();
543544
if (this.history.isAtEnd()) {
544545
this.saveCurrentValue(inputState);
545546
} else {
546-
if (!this.history.has({ text: this._inputEditor.getValue(), state: inputState })) {
547+
const currentEntry = this.getFilteredEntry(this._inputEditor.getValue(), inputState);
548+
if (!this.history.has(currentEntry)) {
547549
this.saveCurrentValue(inputState);
548550
this.history.resetCursor();
549551
}
@@ -552,12 +554,13 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
552554
this.navigateHistory(true);
553555
}
554556

555-
showNextValue(): void {
557+
async showNextValue(): Promise<void> {
556558
const inputState = this.getInputState();
557559
if (this.history.isAtEnd()) {
558560
return;
559561
} else {
560-
if (!this.history.has({ text: this._inputEditor.getValue(), state: inputState })) {
562+
const currentEntry = this.getFilteredEntry(this._inputEditor.getValue(), inputState);
563+
if (!this.history.has(currentEntry)) {
561564
this.saveCurrentValue(inputState);
562565
this.history.resetCursor();
563566
}
@@ -566,11 +569,30 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
566569
this.navigateHistory(false);
567570
}
568571

569-
private navigateHistory(previous: boolean): void {
572+
private async navigateHistory(previous: boolean): Promise<void> {
570573
const historyEntry = previous ?
571574
this.history.previous() : this.history.next();
572575

573-
const historyAttachments = historyEntry.state?.chatContextAttachments ?? [];
576+
let historyAttachments = historyEntry.state?.chatContextAttachments ?? [];
577+
578+
// Check for images in history to restore the value.
579+
if (historyAttachments.length > 0) {
580+
historyAttachments = (await Promise.all(historyAttachments.map(async (attachment) => {
581+
if (attachment.isImage && attachment.references?.length && URI.isUri(attachment.references[0].reference)) {
582+
try {
583+
const buffer = await this.fileService.readFile(attachment.references[0].reference);
584+
const newAttachment = { ...attachment };
585+
newAttachment.value = (isImageVariableEntry(attachment) && attachment.isPasted) ? buffer.value.buffer : await resizeImage(buffer.value.buffer); // if pasted image, we do not need to resize.
586+
return newAttachment;
587+
} catch (err) {
588+
this.logService.error('Failed to restore image from history', err);
589+
return undefined;
590+
}
591+
}
592+
return attachment;
593+
}))).filter(attachment => attachment !== undefined);
594+
}
595+
574596
this._attachmentModel.clearAndSetContext(...historyAttachments);
575597

576598
aria.status(historyEntry.text);
@@ -610,8 +632,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
610632
}
611633

612634
private saveCurrentValue(inputState: IChatInputState): void {
613-
inputState.chatContextAttachments = inputState.chatContextAttachments?.filter(attachment => !attachment.isImage);
614-
const newEntry = { text: this._inputEditor.getValue(), state: inputState };
635+
const newEntry = this.getFilteredEntry(this._inputEditor.getValue(), inputState);
615636
this.history.replaceLast(newEntry);
616637
}
617638

@@ -631,8 +652,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
631652
if (isUserQuery) {
632653
const userQuery = this._inputEditor.getValue();
633654
const inputState = this.getInputState();
634-
inputState.chatContextAttachments = inputState.chatContextAttachments?.filter(attachment => !attachment.isImage);
635-
const entry: IChatHistoryEntry = { text: userQuery, state: inputState };
655+
const entry = this.getFilteredEntry(userQuery, inputState);
636656
this.history.replaceLast(entry);
637657
this.history.add({ text: '' });
638658
}
@@ -648,6 +668,26 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
648668
}
649669
}
650670

671+
// A funtion that filters out specifically the `value` property of the attachment.
672+
private getFilteredEntry(query: string, inputState: IChatInputState): IChatHistoryEntry {
673+
const attachmentsWithoutImageValues = inputState.chatContextAttachments?.map(attachment => {
674+
if (attachment.isImage && attachment.references?.length && attachment.value) {
675+
const newAttachment = { ...attachment };
676+
newAttachment.value = undefined;
677+
return newAttachment;
678+
}
679+
return attachment;
680+
});
681+
682+
inputState.chatContextAttachments = attachmentsWithoutImageValues;
683+
const newEntry = {
684+
text: query,
685+
state: inputState,
686+
};
687+
688+
return newEntry;
689+
}
690+
651691
private _acceptInputForVoiceover(): void {
652692
const domNode = this._inputEditor.getDomNode();
653693
if (!domNode) {

‎src/vs/workbench/contrib/chat/browser/chatPasteProviders.ts

Lines changed: 88 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2,25 +2,28 @@
22
* Copyright (c) Microsoft Corporation. All rights reserved.
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
5-
5+
import { VSBuffer } from '../../../../base/common/buffer.js';
66
import { CancellationToken } from '../../../../base/common/cancellation.js';
7+
import { Codicon } from '../../../../base/common/codicons.js';
78
import { createStringDataTransferItem, IDataTransferItem, IReadonlyVSDataTransfer, VSDataTransfer } from '../../../../base/common/dataTransfer.js';
89
import { HierarchicalKind } from '../../../../base/common/hierarchicalKind.js';
10+
import { Disposable } from '../../../../base/common/lifecycle.js';
11+
import { Mimes } from '../../../../base/common/mime.js';
12+
import { basename, joinPath } from '../../../../base/common/resources.js';
13+
import { URI, UriComponents } from '../../../../base/common/uri.js';
914
import { IRange } from '../../../../editor/common/core/range.js';
1015
import { DocumentPasteContext, DocumentPasteEdit, DocumentPasteEditProvider, DocumentPasteEditsSession } from '../../../../editor/common/languages.js';
1116
import { ITextModel } from '../../../../editor/common/model.js';
1217
import { ILanguageFeaturesService } from '../../../../editor/common/services/languageFeatures.js';
13-
import { Disposable } from '../../../../base/common/lifecycle.js';
14-
import { ChatInputPart } from './chatInputPart.js';
15-
import { IChatWidgetService } from './chat.js';
16-
import { Codicon } from '../../../../base/common/codicons.js';
18+
import { IModelService } from '../../../../editor/common/services/model.js';
1719
import { localize } from '../../../../nls.js';
18-
import { IChatRequestPasteVariableEntry, IChatRequestVariableEntry } from '../common/chatModel.js';
20+
import { IEnvironmentService } from '../../../../platform/environment/common/environment.js';
21+
import { IFileService } from '../../../../platform/files/common/files.js';
22+
import { ILogService } from '../../../../platform/log/common/log.js';
1923
import { IExtensionService, isProposedApiEnabled } from '../../../services/extensions/common/extensions.js';
20-
import { Mimes } from '../../../../base/common/mime.js';
21-
import { IModelService } from '../../../../editor/common/services/model.js';
22-
import { URI, UriComponents } from '../../../../base/common/uri.js';
23-
import { basename } from '../../../../base/common/resources.js';
24+
import { IChatRequestPasteVariableEntry, IChatRequestVariableEntry } from '../common/chatModel.js';
25+
import { IChatWidgetService } from './chat.js';
26+
import { ChatInputPart } from './chatInputPart.js';
2427
import { resizeImage } from './imageUtils.js';
2528

2629
const COPY_MIME_TYPES = 'application/vnd.code.additional-editor-data';
@@ -31,6 +34,7 @@ interface SerializedCopyData {
3134
}
3235

3336
export class PasteImageProvider implements DocumentPasteEditProvider {
37+
private readonly imagesFolder: URI;
3438

3539
public readonly kind = new HierarchicalKind('chat.attach.image');
3640
public readonly providedPasteEditKinds = [this.kind];
@@ -41,7 +45,13 @@ export class PasteImageProvider implements DocumentPasteEditProvider {
4145
constructor(
4246
private readonly chatWidgetService: IChatWidgetService,
4347
private readonly extensionService: IExtensionService,
44-
) { }
48+
@IFileService private readonly fileService: IFileService,
49+
@IEnvironmentService private readonly environmentService: IEnvironmentService,
50+
@ILogService private readonly logService: ILogService,
51+
) {
52+
this.imagesFolder = joinPath(this.environmentService.workspaceStorageHome, 'vscode-chat-images');
53+
this.cleanupOldImages();
54+
}
4555

4656
async provideDocumentPasteEdits(model: ITextModel, ranges: readonly IRange[], dataTransfer: IReadonlyVSDataTransfer, context: DocumentPasteContext, token: CancellationToken): Promise<DocumentPasteEditsSession | undefined> {
4757
if (!this.extensionService.extensions.some(ext => isProposedApiEnabled(ext, 'chatReferenceBinaryData'))) {
@@ -90,12 +100,17 @@ export class PasteImageProvider implements DocumentPasteEditProvider {
90100
tempDisplayName = `${displayName} ${appendValue}`;
91101
}
92102

103+
const fileReference = await this.createFileForMedia(currClipboard, mimeType);
104+
if (token.isCancellationRequested || !fileReference) {
105+
return;
106+
}
107+
93108
const scaledImageData = await resizeImage(currClipboard);
94109
if (token.isCancellationRequested || !scaledImageData) {
95110
return;
96111
}
97112

98-
const scaledImageContext = await getImageAttachContext(scaledImageData, mimeType, token, tempDisplayName);
113+
const scaledImageContext = await getImageAttachContext(scaledImageData, mimeType, token, tempDisplayName, fileReference);
99114
if (token.isCancellationRequested || !scaledImageContext) {
100115
return;
101116
}
@@ -111,21 +126,75 @@ export class PasteImageProvider implements DocumentPasteEditProvider {
111126
const edit = createCustomPasteEdit(model, scaledImageContext, mimeType, this.kind, localize('pastedImageAttachment', 'Pasted Image Attachment'), this.chatWidgetService);
112127
return createEditSession(edit);
113128
}
129+
130+
private async createFileForMedia(
131+
dataTransfer: Uint8Array,
132+
mimeType: string,
133+
): Promise<URI | undefined> {
134+
const exists = await this.fileService.exists(this.imagesFolder);
135+
if (!exists) {
136+
await this.fileService.createFolder(this.imagesFolder);
137+
}
138+
139+
const ext = mimeType.split('/')[1] || 'png';
140+
const filename = `image-${Date.now()}.${ext}`;
141+
const fileUri = joinPath(this.imagesFolder, filename);
142+
143+
const buffer = VSBuffer.wrap(dataTransfer);
144+
await this.fileService.writeFile(fileUri, buffer);
145+
146+
return fileUri;
147+
}
148+
149+
private async cleanupOldImages(): Promise<void> {
150+
const exists = await this.fileService.exists(this.imagesFolder);
151+
if (!exists) {
152+
return;
153+
}
154+
155+
const duration = 7 * 24 * 60 * 60 * 1000; // 7 days
156+
const files = await this.fileService.resolve(this.imagesFolder);
157+
if (!files.children) {
158+
return;
159+
}
160+
161+
await Promise.all(files.children.map(async (file) => {
162+
try {
163+
const timestamp = this.getTimestampFromFilename(file.name);
164+
if (timestamp && (Date.now() - timestamp > duration)) {
165+
await this.fileService.del(file.resource);
166+
}
167+
} catch (err) {
168+
this.logService.error('Failed to clean up old images', err);
169+
}
170+
}));
171+
}
172+
173+
private getTimestampFromFilename(filename: string): number | undefined {
174+
const match = filename.match(/image-(\d+)\./);
175+
if (match) {
176+
return parseInt(match[1], 10);
177+
}
178+
return undefined;
179+
}
114180
}
115181

116-
async function getImageAttachContext(data: Uint8Array, mimeType: string, token: CancellationToken, displayName: string): Promise<IChatRequestVariableEntry | undefined> {
182+
async function getImageAttachContext(data: Uint8Array, mimeType: string, token: CancellationToken, displayName: string, resource: URI): Promise<IChatRequestVariableEntry | undefined> {
117183
const imageHash = await imageToHash(data);
118184
if (token.isCancellationRequested) {
119185
return undefined;
120186
}
121187

122188
return {
189+
kind: 'image',
123190
value: data,
124191
id: imageHash,
125192
name: displayName,
126193
isImage: true,
127194
icon: Codicon.fileMedia,
128-
mimeType
195+
mimeType,
196+
isPasted: true,
197+
references: [{ reference: resource, kind: 'reference' }]
129198
};
130199
}
131200

@@ -308,10 +377,13 @@ export class ChatPasteProvidersFeature extends Disposable {
308377
@ILanguageFeaturesService languageFeaturesService: ILanguageFeaturesService,
309378
@IChatWidgetService chatWidgetService: IChatWidgetService,
310379
@IExtensionService extensionService: IExtensionService,
311-
@IModelService modelService: IModelService
380+
@IFileService fileService: IFileService,
381+
@IModelService modelService: IModelService,
382+
@IEnvironmentService environmentService: IEnvironmentService,
383+
@ILogService logService: ILogService,
312384
) {
313385
super();
314-
this._register(languageFeaturesService.documentPasteEditProvider.register({ scheme: ChatInputPart.INPUT_SCHEME, pattern: '*', hasAccessToAllModels: true }, new PasteImageProvider(chatWidgetService, extensionService)));
386+
this._register(languageFeaturesService.documentPasteEditProvider.register({ scheme: ChatInputPart.INPUT_SCHEME, pattern: '*', hasAccessToAllModels: true }, new PasteImageProvider(chatWidgetService, extensionService, fileService, environmentService, logService)));
315387
this._register(languageFeaturesService.documentPasteEditProvider.register({ scheme: ChatInputPart.INPUT_SCHEME, pattern: '*', hasAccessToAllModels: true }, new PasteTextProvider(chatWidgetService, modelService)));
316388
this._register(languageFeaturesService.documentPasteEditProvider.register('*', new CopyTextProvider()));
317389
}

‎src/vs/workbench/contrib/chat/browser/imageUtils.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ export async function resizeImage(data: Uint8Array | string): Promise<Uint8Array
2727
URL.revokeObjectURL(url);
2828
let { width, height } = img;
2929

30-
if (width < 768 || height < 768) {
30+
if (width <= 768 || height <= 768) {
3131
resolve(data);
3232
return;
3333
}

‎src/vs/workbench/contrib/chat/common/chatModel.ts

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,11 @@ export interface ILinkVariableEntry extends Omit<IBaseChatRequestVariableEntry,
8888
readonly value: URI;
8989
}
9090

91+
export interface IImageVariableEntry extends Omit<IBaseChatRequestVariableEntry, 'kind'> {
92+
readonly kind: 'image';
93+
readonly isPasted?: boolean;
94+
}
95+
9196
export interface IDiagnosticVariableEntryFilterData {
9297
readonly filterUri?: URI;
9398
readonly filterSeverity?: MarkerSeverity;
@@ -120,7 +125,7 @@ export interface IDiagnosticVariableEntry extends Omit<IBaseChatRequestVariableE
120125
readonly kind: 'diagnostic';
121126
}
122127

123-
export type IChatRequestVariableEntry = IChatRequestImplicitVariableEntry | IChatRequestPasteVariableEntry | ISymbolVariableEntry | ICommandResultVariableEntry | ILinkVariableEntry | IBaseChatRequestVariableEntry | IDiagnosticVariableEntry;
128+
export type IChatRequestVariableEntry = IChatRequestImplicitVariableEntry | IChatRequestPasteVariableEntry | ISymbolVariableEntry | ICommandResultVariableEntry | ILinkVariableEntry | IBaseChatRequestVariableEntry | IDiagnosticVariableEntry | IImageVariableEntry;
124129

125130
export function isImplicitVariableEntry(obj: IChatRequestVariableEntry): obj is IChatRequestImplicitVariableEntry {
126131
return obj.kind === 'implicit';
@@ -134,6 +139,10 @@ export function isLinkVariableEntry(obj: IChatRequestVariableEntry): obj is ILin
134139
return obj.kind === 'link';
135140
}
136141

142+
export function isImageVariableEntry(obj: IChatRequestVariableEntry): obj is IImageVariableEntry {
143+
return obj.kind === 'image';
144+
}
145+
137146
export function isDiagnosticsVariableEntry(obj: IChatRequestVariableEntry): obj is IDiagnosticVariableEntry {
138147
return obj.kind === 'diagnostic';
139148
}

0 commit comments

Comments
 (0)