add syntax highlighting support for more languages

vscode/src/autoedits/renderer/decorators/default-decorator.ts

@@ -266,7 +266,7 @@ export class DefaultDecorator implements AutoEditsDecorator {
         const blockifiedAddedLines = blockify(this.editor.document, addedLinesInfo)
         const { dark, light } = generateSuggestionAsImage({
             decorations: blockifiedAddedLines,
-            lang: 'typescript',
+            lang: this.editor.document.languageId,
         })
         const startLineEndColumn = this.getEndColumn(this.editor.document.lineAt(startLine))
 

vscode/src/autoedits/renderer/image-gen/canvas.ts

@@ -1,7 +1,7 @@
 import fs from 'node:fs/promises'
 import path from 'node:path'
 import CanvasKitInit, { type EmulatedCanvas2D } from 'canvaskit-wasm'
-import type { HighlightedAddedLinesDecorationInfo } from './highlight'
+import type { HIGHLIGHT_THEMES, HighlightedAddedLinesDecorationInfo } from './highlight'
 
 type CanvasKitType = Awaited<ReturnType<typeof CanvasKitInit>>
 type RenderContext = {
@@ -76,8 +76,16 @@ function drawText(
     ctx: CanvasRenderingContext2D,
     line: HighlightedAddedLinesDecorationInfo,
     position: { x: number; y: number },
+    mode: keyof typeof HIGHLIGHT_THEMES,
     config: RenderConfig
 ): number {
+    if (!line.highlightedTokens) {
+        // We weren't able to highlight this text, so we just draw the full line
+        ctx.fillStyle = mode === 'dark' ? '#ffffff' : '#000000'
+        ctx.fillText(line.lineText, position.x, position.y + config.fontSize)
+        return ctx.measureText(line.lineText).width
+    }
+
     let xPos = position.x
 
     for (const token of line.highlightedTokens) {
@@ -100,6 +108,12 @@ function drawHighlights(
         return
     }
 
+    if (!line.highlightedTokens) {
+        // No highlighting. TODO Improve this so we still support highlighted regions on non-syntax highlighted
+        // text
+        return
+    }
+
     ctx.fillStyle = config.highlightColor
 
     let xPos = position.x
@@ -131,6 +145,7 @@ function drawHighlights(
 
 export function drawTokensToCanvas(
     highlightedDecorations: HighlightedAddedLinesDecorationInfo[],
+    mode: keyof typeof HIGHLIGHT_THEMES,
     /**
      * Specific configuration to determine how we render the canvas.
      * Consider changing this, or supporting configuration from the user (e.g. font-size)
@@ -161,13 +176,9 @@ export function drawTokensToCanvas(
     // and the required height of the canvas (number of lines determined by their line height)
     let tempYPos = renderConfig.padding.y
     let requiredWidth = 0
-    for (const { highlightedTokens } of highlightedDecorations) {
-        let tempXPos = renderConfig.padding.x
-        for (const token of highlightedTokens) {
-            const measure = tempCtx.measureText(token.content)
-            tempXPos += measure.width
-            requiredWidth = Math.max(requiredWidth, tempXPos)
-        }
+    for (const { lineText } of highlightedDecorations) {
+        const measure = tempCtx.measureText(lineText)
+        requiredWidth = Math.max(requiredWidth, renderConfig.padding.x + measure.width)
         tempYPos += renderConfig.lineHeight
     }
 
@@ -196,7 +207,7 @@ export function drawTokensToCanvas(
         drawHighlights(ctx, line, position, renderConfig)
 
         // Draw text on top
-        drawText(ctx, line, position, renderConfig)
+        drawText(ctx, line, position, mode, renderConfig)
 
         yPos += renderConfig.lineHeight
     }

vscode/src/autoedits/renderer/image-gen/highlight.ts

@@ -1,8 +1,9 @@
 import { type BundledLanguage, type Highlighter, type ThemedToken, createHighlighter } from 'shiki'
+import type { MultiLineSupportedLanguage } from '../../../completions/detect-multiline'
 import type { AddedLinesDecorationInfo } from '../decorators/default-decorator'
 
 export interface HighlightedAddedLinesDecorationInfo extends AddedLinesDecorationInfo {
-    highlightedTokens: ThemedToken[]
+    highlightedTokens?: ThemedToken[]
 }
 
 let highlighter: Highlighter | null = null
@@ -12,15 +13,37 @@ export const HIGHLIGHT_THEMES = {
     dark: 'vitesse-dark',
 } as const
 
-export const HIGHLIGHT_LANGUAGES: Record<string, BundledLanguage> = {
-    TypeScript: 'typescript',
-} as const
+/**
+ * Mapping of supported completiion languages to highlighter languages
+ */
+export const SYNTAX_HIGHLIGHT_MAPPING: Record<MultiLineSupportedLanguage, BundledLanguage> = {
+    astro: 'astro',
+    c: 'c',
+    cpp: 'cpp',
+    csharp: 'csharp',
+    css: 'css',
+    dart: 'dart',
+    elixir: 'elixir',
+    go: 'go',
+    html: 'html',
+    java: 'java',
+    javascript: 'javascript',
+    javascriptreact: 'jsx',
+    kotlin: 'kotlin',
+    php: 'php',
+    python: 'python',
+    rust: 'rust',
+    svelte: 'svelte',
+    typescript: 'typescript',
+    typescriptreact: 'tsx',
+    vue: 'vue',
+}
 
 export async function initHighlighter(): Promise<void> {
     if (!highlighter) {
         highlighter = await createHighlighter({
             themes: Object.values(HIGHLIGHT_THEMES),
-            langs: Object.values(HIGHLIGHT_LANGUAGES),
+            langs: Object.values(SYNTAX_HIGHLIGHT_MAPPING),
         })
     }
 }

vscode/src/autoedits/renderer/image-gen/index.test.ts

@@ -33,27 +33,48 @@ const MOCK_DECORATIONS = [
 
 expect.extend({ toMatchImageSnapshot })
 
-describe('generateSuggestionAsImage', () => {
-    it('generates correct images, with correct highlighting applied, from a set of tokens', async () => {
-        await initImageSuggestionService()
+async function generateImageForTest(
+    decorations: AddedLinesDecorationInfo[],
+    lang: string
+): Promise<{ darkBuffer: Buffer; lightBuffer: Buffer }> {
+    await initImageSuggestionService()
 
-        // These are dataURLs created via .toDataURL('image/png') in CanvasKit.
-        // I need to convert these to images and somehow diff the images with Vitest.
-        // Any ideas would be helpful
-        const { light, dark } = generateSuggestionAsImage({
-            decorations: MOCK_DECORATIONS,
-            lang: 'typescript',
-        })
+    // These are dataURLs created via .toDataURL('image/png') in CanvasKit.
+    // I need to convert these to images and somehow diff the images with Vitest.
+    // Any ideas would be helpful
+    const { light, dark } = generateSuggestionAsImage({
+        decorations,
+        lang,
+    })
 
+    return {
         // These suggestions are generated as dataURLs, so let's convert them back to a useful Buffer for testing
-        const lightBuffer = Buffer.from(light.split(',')[1], 'base64')
-        const darkBuffer = Buffer.from(dark.split(',')[1], 'base64')
+        darkBuffer: Buffer.from(dark.split(',')[1], 'base64'),
+        lightBuffer: Buffer.from(light.split(',')[1], 'base64'),
+    }
+}
+
+describe('generateSuggestionAsImage', () => {
+    it('generates correct images, with correct highlighting applied, from a set of tokens', async () => {
+        const { darkBuffer, lightBuffer } = await generateImageForTest(MOCK_DECORATIONS, 'typescript')
+        expect(lightBuffer).toMatchImageSnapshot({
+            customSnapshotIdentifier: 'highlighted-suggestion-light',
+        })
+        expect(darkBuffer).toMatchImageSnapshot({
+            customSnapshotIdentifier: 'highlighted-suggestion-dark',
+        })
+    })
 
+    it('generates correct images, with correct highlighting applied, from a set of tokens in a language that does not have supported highlighting', async () => {
+        const { darkBuffer, lightBuffer } = await generateImageForTest(
+            MOCK_DECORATIONS,
+            'non-existent-language'
+        )
         expect(lightBuffer).toMatchImageSnapshot({
-            customSnapshotIdentifier: 'generated-suggestion-light',
+            customSnapshotIdentifier: 'unhighlighted-suggestion-light',
         })
         expect(darkBuffer).toMatchImageSnapshot({
-            customSnapshotIdentifier: 'generated-suggestion-dark',
+            customSnapshotIdentifier: 'unhighlighted-suggestion-dark',
         })
     })
 })

vscode/src/autoedits/renderer/image-gen/index.ts

@@ -1,25 +1,30 @@
-import type { BundledLanguage } from 'shiki/langs'
+import type { MultiLineSupportedLanguage } from '../../../completions/detect-multiline'
 import type { AddedLinesDecorationInfo } from '../decorators/default-decorator'
 import { drawTokensToCanvas, initCanvas } from './canvas'
-import { highlightDecorations, initHighlighter } from './highlight'
+import { SYNTAX_HIGHLIGHT_MAPPING, highlightDecorations, initHighlighter } from './highlight'
 
 export async function initImageSuggestionService() {
     return Promise.all([initHighlighter(), initCanvas()])
 }
 
 interface SuggestionOptions {
     decorations: AddedLinesDecorationInfo[]
-    lang: BundledLanguage
+    lang: string
 }
 
 export function generateSuggestionAsImage(options: SuggestionOptions): { light: string; dark: string } {
     const { decorations, lang } = options
+    const highlightingLang = SYNTAX_HIGHLIGHT_MAPPING[lang as MultiLineSupportedLanguage]
 
-    const highlightedLightDecorations = highlightDecorations(decorations, lang, 'light')
-    const highlightedDarkDecorations = highlightDecorations(decorations, lang, 'dark')
+    const darkDecorations = highlightingLang
+        ? highlightDecorations(decorations, highlightingLang, 'dark')
+        : decorations
+    const lightDecorations = highlightingLang
+        ? highlightDecorations(decorations, highlightingLang, 'light')
+        : decorations
 
     return {
-        dark: drawTokensToCanvas(highlightedDarkDecorations).toDataURL('image/png'),
-        light: drawTokensToCanvas(highlightedLightDecorations).toDataURL('image/png'),
+        dark: drawTokensToCanvas(darkDecorations, 'dark').toDataURL('image/png'),
+        light: drawTokensToCanvas(lightDecorations, 'light').toDataURL('image/png'),
     }
 }

vscode/src/completions/detect-multiline.ts

@@ -31,7 +31,7 @@ export function endsWithBlockStart(text: string, languageId: string): string | n
 
 // Languages with more than 100 multiline completions in the last month and CAR > 20%:
 // https://sourcegraph.looker.com/explore/sourcegraph/cody?qid=JBItVt6VFMlCtMa9KOBmjh&origin_space=562
-const LANGUAGES_WITH_MULTILINE_SUPPORT = [
+export const LANGUAGES_WITH_MULTILINE_SUPPORT = [
     'astro',
     'c',
     'cpp',
@@ -52,13 +52,17 @@ const LANGUAGES_WITH_MULTILINE_SUPPORT = [
     'typescript',
     'typescriptreact',
     'vue',
-]
+] as const
+
+export type MultiLineSupportedLanguage = (typeof LANGUAGES_WITH_MULTILINE_SUPPORT)[number]
 
 export function detectMultiline(params: DetectMultilineParams): DetectMultilineResult {
     const { docContext, languageId, position } = params
     const { prefix, prevNonEmptyLine, nextNonEmptyLine, currentLinePrefix, currentLineSuffix } =
         docContext
-    const isMultilineSupported = LANGUAGES_WITH_MULTILINE_SUPPORT.includes(languageId)
+    const isMultilineSupported = LANGUAGES_WITH_MULTILINE_SUPPORT.includes(
+        languageId as MultiLineSupportedLanguage
+    )
 
     const blockStart = endsWithBlockStart(prefix, languageId)
     const isBlockStartActive = Boolean(blockStart)