Skip to content

Commit

Permalink
adds images in history, create file for pasted images in chat (#241664)
Browse files Browse the repository at this point in the history 
* 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
  • Loading branch information
@justschen
justschen authored Feb 24, 2025
1 parent f0df737 commit ac0e2ac
Show file tree
Hide file tree
Showing 4 changed files with 150 additions and 29 deletions.
62 changes: 51 additions & 11 deletions src/vs/workbench/contrib/chat/browser/chatInputPart.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ import { revealInSideBarCommand } from '../../files/browser/fileActions.contribu
import { ChatAgentLocation, IChatAgentService } from '../common/chatAgents.js';
import { ChatContextKeys } from '../common/chatContextKeys.js';
import { IChatEditingSession, WorkingSetEntryRemovalReason, WorkingSetEntryState } from '../common/chatEditingService.js';
import { IChatRequestVariableEntry, isLinkVariableEntry, isPasteVariableEntry } from '../common/chatModel.js';
import { IChatRequestVariableEntry, isImageVariableEntry, isLinkVariableEntry, isPasteVariableEntry } from '../common/chatModel.js';
import { IChatFollowup } from '../common/chatService.js';
import { IChatVariablesService } from '../common/chatVariables.js';
import { IChatResponseViewModel } from '../common/chatViewModel.js';
Expand All @@ -101,6 +101,7 @@ import { ChatFollowups } from './chatFollowups.js';
import { IChatViewState } from './chatWidget.js';
import { ChatFileReference } from './contrib/chatDynamicVariables/chatFileReference.js';
import { ChatImplicitContext } from './contrib/chatImplicitContext.js';
import { resizeImage } from './imageUtils.js';

const $ = dom.$;

Expand Down Expand Up @@ -538,12 +539,13 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
return this.container;
}

showPreviousValue(): void {
async showPreviousValue(): Promise<void> {
const inputState = this.getInputState();
if (this.history.isAtEnd()) {
this.saveCurrentValue(inputState);
} else {
if (!this.history.has({ text: this._inputEditor.getValue(), state: inputState })) {
const currentEntry = this.getFilteredEntry(this._inputEditor.getValue(), inputState);
if (!this.history.has(currentEntry)) {
this.saveCurrentValue(inputState);
this.history.resetCursor();
}
Expand All @@ -552,12 +554,13 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
this.navigateHistory(true);
}

showNextValue(): void {
async showNextValue(): Promise<void> {
const inputState = this.getInputState();
if (this.history.isAtEnd()) {
return;
} else {
if (!this.history.has({ text: this._inputEditor.getValue(), state: inputState })) {
const currentEntry = this.getFilteredEntry(this._inputEditor.getValue(), inputState);
if (!this.history.has(currentEntry)) {
this.saveCurrentValue(inputState);
this.history.resetCursor();
}
Expand All @@ -566,11 +569,30 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
this.navigateHistory(false);
}

private navigateHistory(previous: boolean): void {
private async navigateHistory(previous: boolean): Promise<void> {
const historyEntry = previous ?
this.history.previous() : this.history.next();

const historyAttachments = historyEntry.state?.chatContextAttachments ?? [];
let historyAttachments = historyEntry.state?.chatContextAttachments ?? [];

// Check for images in history to restore the value.
if (historyAttachments.length > 0) {
historyAttachments = (await Promise.all(historyAttachments.map(async (attachment) => {
if (attachment.isImage && attachment.references?.length && URI.isUri(attachment.references[0].reference)) {
try {
const buffer = await this.fileService.readFile(attachment.references[0].reference);
const newAttachment = { ...attachment };
newAttachment.value = (isImageVariableEntry(attachment) && attachment.isPasted) ? buffer.value.buffer : await resizeImage(buffer.value.buffer); // if pasted image, we do not need to resize.
return newAttachment;
} catch (err) {
this.logService.error('Failed to restore image from history', err);
return undefined;
}
}
return attachment;
}))).filter(attachment => attachment !== undefined);
}

this._attachmentModel.clearAndSetContext(...historyAttachments);

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

private saveCurrentValue(inputState: IChatInputState): void {
inputState.chatContextAttachments = inputState.chatContextAttachments?.filter(attachment => !attachment.isImage);
const newEntry = { text: this._inputEditor.getValue(), state: inputState };
const newEntry = this.getFilteredEntry(this._inputEditor.getValue(), inputState);
this.history.replaceLast(newEntry);
}

Expand All @@ -631,8 +652,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
if (isUserQuery) {
const userQuery = this._inputEditor.getValue();
const inputState = this.getInputState();
inputState.chatContextAttachments = inputState.chatContextAttachments?.filter(attachment => !attachment.isImage);
const entry: IChatHistoryEntry = { text: userQuery, state: inputState };
const entry = this.getFilteredEntry(userQuery, inputState);
this.history.replaceLast(entry);
this.history.add({ text: '' });
}
Expand All @@ -648,6 +668,26 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
}
}

// A funtion that filters out specifically the `value` property of the attachment.
private getFilteredEntry(query: string, inputState: IChatInputState): IChatHistoryEntry {
const attachmentsWithoutImageValues = inputState.chatContextAttachments?.map(attachment => {
if (attachment.isImage && attachment.references?.length && attachment.value) {
const newAttachment = { ...attachment };
newAttachment.value = undefined;
return newAttachment;
}
return attachment;
});

inputState.chatContextAttachments = attachmentsWithoutImageValues;
const newEntry = {
text: query,
state: inputState,
};

return newEntry;
}

private _acceptInputForVoiceover(): void {
const domNode = this._inputEditor.getDomNode();
if (!domNode) {
Expand Down
104 changes: 88 additions & 16 deletions src/vs/workbench/contrib/chat/browser/chatPasteProviders.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,25 +2,28 @@
* Copyright (c) Microsoft Corporation. All rights reserved.
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/

import { VSBuffer } from '../../../../base/common/buffer.js';
import { CancellationToken } from '../../../../base/common/cancellation.js';
import { Codicon } from '../../../../base/common/codicons.js';
import { createStringDataTransferItem, IDataTransferItem, IReadonlyVSDataTransfer, VSDataTransfer } from '../../../../base/common/dataTransfer.js';
import { HierarchicalKind } from '../../../../base/common/hierarchicalKind.js';
import { Disposable } from '../../../../base/common/lifecycle.js';
import { Mimes } from '../../../../base/common/mime.js';
import { basename, joinPath } from '../../../../base/common/resources.js';
import { URI, UriComponents } from '../../../../base/common/uri.js';
import { IRange } from '../../../../editor/common/core/range.js';
import { DocumentPasteContext, DocumentPasteEdit, DocumentPasteEditProvider, DocumentPasteEditsSession } from '../../../../editor/common/languages.js';
import { ITextModel } from '../../../../editor/common/model.js';
import { ILanguageFeaturesService } from '../../../../editor/common/services/languageFeatures.js';
import { Disposable } from '../../../../base/common/lifecycle.js';
import { ChatInputPart } from './chatInputPart.js';
import { IChatWidgetService } from './chat.js';
import { Codicon } from '../../../../base/common/codicons.js';
import { IModelService } from '../../../../editor/common/services/model.js';
import { localize } from '../../../../nls.js';
import { IChatRequestPasteVariableEntry, IChatRequestVariableEntry } from '../common/chatModel.js';
import { IEnvironmentService } from '../../../../platform/environment/common/environment.js';
import { IFileService } from '../../../../platform/files/common/files.js';
import { ILogService } from '../../../../platform/log/common/log.js';
import { IExtensionService, isProposedApiEnabled } from '../../../services/extensions/common/extensions.js';
import { Mimes } from '../../../../base/common/mime.js';
import { IModelService } from '../../../../editor/common/services/model.js';
import { URI, UriComponents } from '../../../../base/common/uri.js';
import { basename } from '../../../../base/common/resources.js';
import { IChatRequestPasteVariableEntry, IChatRequestVariableEntry } from '../common/chatModel.js';
import { IChatWidgetService } from './chat.js';
import { ChatInputPart } from './chatInputPart.js';
import { resizeImage } from './imageUtils.js';

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

export class PasteImageProvider implements DocumentPasteEditProvider {
private readonly imagesFolder: URI;

public readonly kind = new HierarchicalKind('chat.attach.image');
public readonly providedPasteEditKinds = [this.kind];
Expand All @@ -41,7 +45,13 @@ export class PasteImageProvider implements DocumentPasteEditProvider {
constructor(
private readonly chatWidgetService: IChatWidgetService,
private readonly extensionService: IExtensionService,
) { }
@IFileService private readonly fileService: IFileService,
@IEnvironmentService private readonly environmentService: IEnvironmentService,
@ILogService private readonly logService: ILogService,
) {
this.imagesFolder = joinPath(this.environmentService.workspaceStorageHome, 'vscode-chat-images');
this.cleanupOldImages();
}

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

const fileReference = await this.createFileForMedia(currClipboard, mimeType);
if (token.isCancellationRequested || !fileReference) {
return;
}

const scaledImageData = await resizeImage(currClipboard);
if (token.isCancellationRequested || !scaledImageData) {
return;
}

const scaledImageContext = await getImageAttachContext(scaledImageData, mimeType, token, tempDisplayName);
const scaledImageContext = await getImageAttachContext(scaledImageData, mimeType, token, tempDisplayName, fileReference);
if (token.isCancellationRequested || !scaledImageContext) {
return;
}
Expand All @@ -111,21 +126,75 @@ export class PasteImageProvider implements DocumentPasteEditProvider {
const edit = createCustomPasteEdit(model, scaledImageContext, mimeType, this.kind, localize('pastedImageAttachment', 'Pasted Image Attachment'), this.chatWidgetService);
return createEditSession(edit);
}

private async createFileForMedia(
dataTransfer: Uint8Array,
mimeType: string,
): Promise<URI | undefined> {
const exists = await this.fileService.exists(this.imagesFolder);
if (!exists) {
await this.fileService.createFolder(this.imagesFolder);
}

const ext = mimeType.split('/')[1] || 'png';
const filename = `image-${Date.now()}.${ext}`;
const fileUri = joinPath(this.imagesFolder, filename);

const buffer = VSBuffer.wrap(dataTransfer);
await this.fileService.writeFile(fileUri, buffer);

return fileUri;
}

private async cleanupOldImages(): Promise<void> {
const exists = await this.fileService.exists(this.imagesFolder);
if (!exists) {
return;
}

const duration = 7 * 24 * 60 * 60 * 1000; // 7 days
const files = await this.fileService.resolve(this.imagesFolder);
if (!files.children) {
return;
}

await Promise.all(files.children.map(async (file) => {
try {
const timestamp = this.getTimestampFromFilename(file.name);
if (timestamp && (Date.now() - timestamp > duration)) {
await this.fileService.del(file.resource);
}
} catch (err) {
this.logService.error('Failed to clean up old images', err);
}
}));
}

private getTimestampFromFilename(filename: string): number | undefined {
const match = filename.match(/image-(\d+)\./);
if (match) {
return parseInt(match[1], 10);
}
return undefined;
}
}

async function getImageAttachContext(data: Uint8Array, mimeType: string, token: CancellationToken, displayName: string): Promise<IChatRequestVariableEntry | undefined> {
async function getImageAttachContext(data: Uint8Array, mimeType: string, token: CancellationToken, displayName: string, resource: URI): Promise<IChatRequestVariableEntry | undefined> {
const imageHash = await imageToHash(data);
if (token.isCancellationRequested) {
return undefined;
}

return {
kind: 'image',
value: data,
id: imageHash,
name: displayName,
isImage: true,
icon: Codicon.fileMedia,
mimeType
mimeType,
isPasted: true,
references: [{ reference: resource, kind: 'reference' }]
};
}

Expand Down Expand Up @@ -308,10 +377,13 @@ export class ChatPasteProvidersFeature extends Disposable {
@ILanguageFeaturesService languageFeaturesService: ILanguageFeaturesService,
@IChatWidgetService chatWidgetService: IChatWidgetService,
@IExtensionService extensionService: IExtensionService,
@IModelService modelService: IModelService
@IFileService fileService: IFileService,
@IModelService modelService: IModelService,
@IEnvironmentService environmentService: IEnvironmentService,
@ILogService logService: ILogService,
) {
super();
this._register(languageFeaturesService.documentPasteEditProvider.register({ scheme: ChatInputPart.INPUT_SCHEME, pattern: '*', hasAccessToAllModels: true }, new PasteImageProvider(chatWidgetService, extensionService)));
this._register(languageFeaturesService.documentPasteEditProvider.register({ scheme: ChatInputPart.INPUT_SCHEME, pattern: '*', hasAccessToAllModels: true }, new PasteImageProvider(chatWidgetService, extensionService, fileService, environmentService, logService)));
this._register(languageFeaturesService.documentPasteEditProvider.register({ scheme: ChatInputPart.INPUT_SCHEME, pattern: '*', hasAccessToAllModels: true }, new PasteTextProvider(chatWidgetService, modelService)));
this._register(languageFeaturesService.documentPasteEditProvider.register('*', new CopyTextProvider()));
}
Expand Down
2 changes: 1 addition & 1 deletion src/vs/workbench/contrib/chat/browser/imageUtils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ export async function resizeImage(data: Uint8Array | string): Promise<Uint8Array
URL.revokeObjectURL(url);
let { width, height } = img;

if (width < 768 || height < 768) {
if (width <= 768 || height <= 768) {
resolve(data);
return;
}
Expand Down
11 changes: 10 additions & 1 deletion src/vs/workbench/contrib/chat/common/chatModel.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,11 @@ export interface ILinkVariableEntry extends Omit<IBaseChatRequestVariableEntry,
readonly value: URI;
}

export interface IImageVariableEntry extends Omit<IBaseChatRequestVariableEntry, 'kind'> {
readonly kind: 'image';
readonly isPasted?: boolean;
}

export interface IDiagnosticVariableEntryFilterData {
readonly filterUri?: URI;
readonly filterSeverity?: MarkerSeverity;
Expand Down Expand Up @@ -120,7 +125,7 @@ export interface IDiagnosticVariableEntry extends Omit<IBaseChatRequestVariableE
readonly kind: 'diagnostic';
}

export type IChatRequestVariableEntry = IChatRequestImplicitVariableEntry | IChatRequestPasteVariableEntry | ISymbolVariableEntry | ICommandResultVariableEntry | ILinkVariableEntry | IBaseChatRequestVariableEntry | IDiagnosticVariableEntry;
export type IChatRequestVariableEntry = IChatRequestImplicitVariableEntry | IChatRequestPasteVariableEntry | ISymbolVariableEntry | ICommandResultVariableEntry | ILinkVariableEntry | IBaseChatRequestVariableEntry | IDiagnosticVariableEntry | IImageVariableEntry;

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

export function isImageVariableEntry(obj: IChatRequestVariableEntry): obj is IImageVariableEntry {
return obj.kind === 'image';
}

export function isDiagnosticsVariableEntry(obj: IChatRequestVariableEntry): obj is IDiagnosticVariableEntry {
return obj.kind === 'diagnostic';
}
Expand Down

0 comments on commit ac0e2ac

Please sign in to comment.