Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

adds images in history, create file for pasted images in chat #241664

Merged
merged 7 commits into from
Feb 24, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 51 additions & 11 deletions src/vs/workbench/contrib/chat/browser/chatInputPart.ts
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,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 @@ -100,6 +100,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 @@ -535,12 +536,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 @@ -549,12 +551,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 @@ -563,11 +566,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 @@ -607,8 +629,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 @@ -628,8 +649,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 @@ -645,6 +665,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) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there some case where an image attachment doesn't have references, and then the image data will still be saved in history?

Copy link
Contributor Author

@justschen justschen Feb 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

there shouldn't be case where images don't have references (besides if it comes from web and is an image, then it is okay because the value will be a URL)

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
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
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
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
Loading