Fix vuln logging in code blocks

Author: mjbvz

src/vs/workbench/contrib/chat/browser/codeBlockPart.ts

@@ -320,6 +320,7 @@ export class CodeBlockPart extends Disposable {
 	}
 
 	async render(data: ICodeBlockData, width: number, editable: boolean | undefined) {
+		this.currentCodeBlockData = data;
 		if (data.parentContextKeyService) {
 			this.contextKeyService.updateParent(data.parentContextKeyService);
 		}

src/vs/workbench/contrib/chat/common/annotations.ts

@@ -7,7 +7,7 @@ import { basename } from 'vs/base/common/resources';
 import { URI } from 'vs/base/common/uri';
 import { IRange } from 'vs/editor/common/core/range';
 import { IChatProgressRenderableResponseContent, IChatProgressResponseContent } from 'vs/workbench/contrib/chat/common/chatModel';
-import { IChatAgentMarkdownContentWithVulnerability, IChatAgentVulnerabilityDetails, IChatContentInlineReference } from 'vs/workbench/contrib/chat/common/chatService';
+import { IChatAgentMarkdownContentWithVulnerability, IChatAgentVulnerabilityDetails, IChatContentInlineReference, IChatMarkdownContent } from 'vs/workbench/contrib/chat/common/chatService';
 
 export const contentRefUrl = 'http://_vscodecontentref_'; // must be lowercase for URI
 
@@ -43,16 +43,40 @@ export function annotateSpecialMarkdownContent(response: ReadonlyArray<IChatProg
 }
 
 export interface IMarkdownVulnerability {
-	title: string;
-	description: string;
-	range: IRange;
+	readonly title: string;
+	readonly description: string;
+	readonly range: IRange;
+}
+
+export function annotateVulnerabilitiesInText(response: ReadonlyArray<IChatProgressResponseContent>): readonly IChatMarkdownContent[] {
+	const result: IChatMarkdownContent[] = [];
+	for (const item of response) {
+		const previousItem = result[result.length - 1];
+		if (item.kind === 'markdownContent') {
+			if (previousItem?.kind === 'markdownContent') {
+				result[result.length - 1] = { content: new MarkdownString(previousItem.content.value + item.content.value, { isTrusted: previousItem.content.isTrusted }), kind: 'markdownContent' };
+			} else {
+				result.push(item);
+			}
+		} else if (item.kind === 'markdownVuln') {
+			const vulnText = encodeURIComponent(JSON.stringify(item.vulnerabilities));
+			const markdownText = `<vscode_annotation details='${vulnText}'>${item.content.value}</vscode_annotation>`;
+			if (previousItem?.kind === 'markdownContent') {
+				result[result.length - 1] = { content: new MarkdownString(previousItem.content.value + markdownText, { isTrusted: previousItem.content.isTrusted }), kind: 'markdownContent' };
+			} else {
+				result.push({ content: new MarkdownString(markdownText), kind: 'markdownContent' });
+			}
+		}
+	}
+
+	return result;
 }
 
 export function extractVulnerabilitiesFromText(text: string): { newText: string; vulnerabilities: IMarkdownVulnerability[] } {
 	const vulnerabilities: IMarkdownVulnerability[] = [];
 	let newText = text;
 	let match: RegExpExecArray | null;
-	while ((match = /<vscode_annotation details="(.*?)">(.*?)<\/vscode_annotation>/ms.exec(newText)) !== null) {
+	while ((match = /<vscode_annotation details=["'](.*?)["']>(.*?)<\/vscode_annotation>/ms.exec(newText)) !== null) {
 		const [full, details, content] = match;
 		const start = match.index;
 		const textBefore = newText.substring(0, start);

src/vs/workbench/contrib/chat/common/chatViewModel.ts

@@ -5,19 +5,20 @@
 
 import { Emitter, Event } from 'vs/base/common/event';
 import { Disposable } from 'vs/base/common/lifecycle';
+import { ResourceMap } from 'vs/base/common/map';
 import { marked } from 'vs/base/common/marked/marked';
 import { ThemeIcon } from 'vs/base/common/themables';
 import { URI } from 'vs/base/common/uri';
+import { TextEdit } from 'vs/editor/common/languages';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { ILogService } from 'vs/platform/log/common/log';
+import { annotateVulnerabilitiesInText } from 'vs/workbench/contrib/chat/common/annotations';
 import { IChatAgentCommand, IChatAgentData, IChatAgentResult } from 'vs/workbench/contrib/chat/common/chatAgents';
 import { ChatModelInitState, IChatModel, IChatRequestModel, IChatResponseModel, IChatWelcomeMessageContent, IResponse } from 'vs/workbench/contrib/chat/common/chatModel';
 import { IParsedChatRequest } from 'vs/workbench/contrib/chat/common/chatParserTypes';
 import { IChatCommandButton, IChatContentReference, IChatFollowup, IChatProgressMessage, IChatResponseErrorDetails, IChatResponseProgressFileTreeData, IChatUsedContext, InteractiveSessionVoteDirection } from 'vs/workbench/contrib/chat/common/chatService';
 import { countWords } from 'vs/workbench/contrib/chat/common/chatWordCounter';
 import { CodeBlockModelCollection } from './codeBlockModelCollection';
-import { TextEdit } from 'vs/editor/common/languages';
-import { ResourceMap } from 'vs/base/common/map';
 
 export function isRequestVM(item: unknown): item is IChatRequestViewModel {
 	return !!item && typeof item === 'object' && 'message' in item;
@@ -259,10 +260,15 @@ export class ChatViewModel extends Disposable implements IChatViewModel {
 	}
 
 	updateCodeBlockTextModels(model: IChatRequestViewModel | IChatResponseViewModel) {
-		const content = isRequestVM(model) ? model.messageText : model.response.asString();
-		const renderer = new marked.Renderer();
+		let content: string;
+		if (isRequestVM(model)) {
+			content = model.messageText;
+		} else {
+			content = annotateVulnerabilitiesInText(model.response.value).map(x => x.content.value).join('');
+		}
 
 		let codeBlockIndex = 0;
+		const renderer = new marked.Renderer();
 		renderer.code = (value, languageId) => {
 			languageId ??= '';
 			const newText = this.fixCodeText(value, languageId);

src/vs/workbench/contrib/chat/common/codeBlockModelCollection.ts

@@ -34,16 +34,16 @@ export class CodeBlockModelCollection extends Disposable {
 		this.clear();
 	}
 
-	get(sessionId: string, chat: IChatRequestViewModel | IChatResponseViewModel, codeBlockIndex: number): { model: Promise<IResolvedTextEditorModel>; vulns: readonly IMarkdownVulnerability[] } | undefined {
+	get(sessionId: string, chat: IChatRequestViewModel | IChatResponseViewModel, codeBlockIndex: number): { model: Promise<IResolvedTextEditorModel>; readonly vulns: readonly IMarkdownVulnerability[] } | undefined {
 		const uri = this.getUri(sessionId, chat, codeBlockIndex);
-		const entries = this._models.get(uri);
-		if (!entries) {
+		const entry = this._models.get(uri);
+		if (!entry) {
 			return;
 		}
-		return { model: entries.model.then(ref => ref.object), vulns: entries.vulns };
+		return { model: entry.model.then(ref => ref.object), vulns: entry.vulns };
 	}
 
-	getOrCreate(sessionId: string, chat: IChatRequestViewModel | IChatResponseViewModel, codeBlockIndex: number): { model: Promise<IResolvedTextEditorModel>; vulns: readonly IMarkdownVulnerability[] } {
+	getOrCreate(sessionId: string, chat: IChatRequestViewModel | IChatResponseViewModel, codeBlockIndex: number): { model: Promise<IResolvedTextEditorModel>; readonly vulns: readonly IMarkdownVulnerability[] } {
 		const existing = this.get(sessionId, chat, codeBlockIndex);
 		if (existing) {
 			return existing;
@@ -65,7 +65,7 @@ export class CodeBlockModelCollection extends Disposable {
 
 		const extractedVulns = extractVulnerabilitiesFromText(content.text);
 		const newText = extractedVulns.newText;
-		entry.vulns = extractedVulns.vulnerabilities;
+		this.setVulns(sessionId, chat, codeBlockIndex, extractedVulns.vulnerabilities);
 
 		const textModel = (await entry.model).textEditorModel;
 		if (content.languageId) {
@@ -91,6 +91,14 @@ export class CodeBlockModelCollection extends Disposable {
 		}
 	}
 
+	private setVulns(sessionId: string, chat: IChatRequestViewModel | IChatResponseViewModel, codeBlockIndex: number, vulnerabilities: IMarkdownVulnerability[]) {
+		const uri = this.getUri(sessionId, chat, codeBlockIndex);
+		const entry = this._models.get(uri);
+		if (entry) {
+			entry.vulns = vulnerabilities;
+		}
+	}
+
 	private getUri(sessionId: string, chat: IChatRequestViewModel | IChatResponseViewModel, index: number): URI {
 		const metadata = this.getUriMetaData(chat);
 		return URI.from({