Build an AI-Powered VS Code Extension: Auto-Generate Comments and Documentation

This article contains affiliate links. I may earn a commission at no extra cost to you.

---

title: "Build an AI-Powered VS Code Extension: Auto-Generate Comments and Documentation"
published: true
description: "Learn to create a VS Code extension that uses AI to automatically generate intelligent comments and documentation for your code"
tags: ai, vscode, tutorial, automation, productivity

cover_image: https://dev-to-uploads.s3.amazonaws.com/uploads/articles/vscode-ai-extension-cover.png

Writing good documentation is crucial, but let's be honest—it's often the last thing we want to do after solving a complex problem. What if your editor could help by analyzing your code and suggesting meaningful comments automatically?

In this tutorial, we'll build a VS Code extension that integrates with OpenAI's API to generate intelligent comments and documentation. You'll learn the fundamentals of VS Code extension development while creating something genuinely useful for your daily workflow.

Prerequisites

Before we start, make sure you have:

• Node.js (v16 or higher)
• VS Code installed
• An OpenAI API key
• Basic familiarity with TypeScript/JavaScript

Setting Up the Development Environment

First, let's install the necessary tools and create our extension scaffold:

npm install -g yo generator-code
yo code

When prompted, choose:

• "New Extension (TypeScript)"
• Name: "ai-code-documenter"
• Identifier: Accept default
• Description: "AI-powered code documentation generator"
• Initialize git repository: Yes
• Bundle source code: No
• Package manager: npm

This creates a basic extension structure. Navigate to the new directory:

cd ai-code-documenter
npm install

Project Structure Overview

Your extension now has this structure:

ai-code-documenter/
├── src/
│   └── extension.ts     # Main extension logic
├── package.json         # Extension manifest
└── tsconfig.json       # TypeScript config

The package.json file is particularly important—it defines your extension's capabilities, commands, and configuration options.

Installing Dependencies

We'll need the OpenAI SDK for API integration:

npm install openai
npm install --save-dev @types/vscode

Configuring the Extension Manifest

Update your package.json to define commands and configuration options:

{
  "contributes": {
    "commands": [
      {
        "command": "ai-documenter.generateComment",
        "title": "Generate AI Comment",
        "category": "AI Documenter"
      },
      {
        "command": "ai-documenter.documentFunction",
        "title": "Document Function",
        "category": "AI Documenter"
      }
    ],
    "menus": {
      "editor/context": [
        {
          "command": "ai-documenter.generateComment",
          "when": "editorHasSelection",
          "group": "ai-documenter"
        }
      ]
    },
    "configuration": {
      "title": "AI Code Documenter",
      "properties": {
        "aiDocumenter.apiKey": {
          "type": "string",
          "default": "",
          "description": "OpenAI API Key",
          "scope": "application"
        },
        "aiDocumenter.model": {
          "type": "string",
          "default": "gpt-3.5-turbo",
          "enum": ["gpt-3.5-turbo", "gpt-4"],
          "description": "AI model to use for generation"
        },
        "aiDocumenter.commentStyle": {
          "type": "string",
          "default": "jsdoc",
          "enum": ["jsdoc", "inline", "block"],
          "description": "Comment style preference"
        }
      }
    }
  }
}

Building the Core Extension Logic

Now let's implement the main functionality in src/extension.ts:

import * as vscode from 'vscode';
import OpenAI from 'openai';

let openai: OpenAI | null = null;

export function activate(context: vscode.ExtensionContext) {
    console.log('AI Code Documenter is now active!');

    // Initialize OpenAI client
    initializeOpenAI();

    // Register commands
    const generateComment = vscode.commands.registerCommand(
        'ai-documenter.generateComment',
        generateCommentCommand
    );

    const documentFunction = vscode.commands.registerCommand(
        'ai-documenter.documentFunction',
        documentFunctionCommand
    );

    context.subscriptions.push(generateComment, documentFunction);

    // Watch for configuration changes
    vscode.workspace.onDidChangeConfiguration(event => {
        if (event.affectsConfiguration('aiDocumenter.apiKey')) {
            initializeOpenAI();
        }
    });
}

function initializeOpenAI() {
    const config = vscode.workspace.getConfiguration('aiDocumenter');
    const apiKey = config.get<string>('apiKey');

    if (apiKey) {
        openai = new OpenAI({ apiKey });
    } else {
        openai = null;
        vscode.window.showWarningMessage(
            'Please set your OpenAI API key in settings to use AI Documenter'
        );
    }
}

async function generateCommentCommand() {
    if (!openai) {
        vscode.window.showErrorMessage('OpenAI API key not configured');
        return;
    }

    const editor = vscode.window.activeTextEditor;
    if (!editor) {
        return;
    }

    const selection = editor.selection;
    const selectedText = editor.document.getText(selection);

    if (!selectedText) {
        vscode.window.showInformationMessage('Please select code to document');
        return;
    }

    try {
        vscode.window.withProgress({
            location: vscode.ProgressLocation.Notification,
            title: "Generating comment...",
            cancellable: false
        }, async () => {
            const comment = await generateComment(selectedText);
            await insertComment(editor, selection, comment);
        });
    } catch (error) {
        vscode.window.showErrorMessage(`Error generating comment: ${error}`);
    }
}

async function generateComment(code: string): Promise<string> {
    const config = vscode.workspace.getConfiguration('aiDocumenter');
    const model = config.get<string>('model', 'gpt-3.5-turbo');
    const commentStyle = config.get<string>('commentStyle', 'jsdoc');

    const prompt = createPrompt(code, commentStyle);

    const response = await openai!.chat.completions.create({
        model,
        messages: [{ role: 'user', content: prompt }],
        max_tokens: 200,
        temperature: 0.3
    });

    return response.choices[0]?.message?.content?.trim() || '';
}

function createPrompt(code: string, style: string): string {
    const styleInstructions = {
        jsdoc: 'Generate a JSDoc-style comment',
        inline: 'Generate a concise inline comment',
        block: 'Generate a multi-line block comment'
    };

    return `${styleInstructions[style as keyof typeof styleInstructions]} for this code. Focus on explaining what the code does, its parameters, and return value if applicable. Be concise and helpful:\n\n${code}`;
}

async function insertComment(editor: vscode.TextEditor, selection: vscode.Selection, comment: string) {
    const line = selection.start.line;
    const insertPosition = new vscode.Position(line, 0);

    await editor.edit(editBuilder => {
        editBuilder.insert(insertPosition, comment + '\n');
    });
}

async function documentFunctionCommand() {
    const editor = vscode.window.activeTextEditor;
    if (!editor) {
        return;
    }

    const position = editor.selection.active;
    const functionCode = extractFunctionAtPosition(editor.document, position);

    if (!functionCode) {
        vscode.window.showInformationMessage('No function found at cursor position');
        return;
    }

    try {
        const comment = await generateComment(functionCode);
        const functionStart = findFunctionStart(editor.document, position);

        if (functionStart) {
            await editor.edit(editBuilder => {
                editBuilder.insert(functionStart, comment + '\n');
            });
        }
    } catch (error) {
        vscode.window.showErrorMessage(`Error documenting function: ${error}`);
    }
}

function extractFunctionAtPosition(document: vscode.TextDocument, position: vscode.Position): string | null {
    // Simple function extraction - you could make this more sophisticated
    const line = document.lineAt(position);
    const text = line.text;

    // Look for function declarations
    if (text.includes('function') || text.includes('=>') || text.includes('def ')) {
        let startLine = position.line;
        let endLine = position.line;

        // Find function end (simplified)
        while (endLine < document.lineCount - 1) {
            const nextLine = document.lineAt(endLine + 1);
            if (nextLine.text.trim() === '' && document.lineAt(endLine).text.includes('}')) {
                break;
            }
            endLine++;
        }

        const range = new vscode.Range(startLine, 0, endLine, document.lineAt(endLine).text.length);
        return document.getText(range);
    }

    return null;
}

function findFunctionStart(document: vscode.TextDocument, position: vscode.Position): vscode.Position | null {
    // Find the start of the function to insert comment above it
    for (let i = position.line; i >= 0; i--) {
        const line = document.lineAt(i);
        if (line.text.includes('function') || line.text.includes('=>') || line.text.includes('def ')) {
            return new vscode.Position(i, 0);
        }
    }
    return null;
}

export function deactivate() {}

Testing Your Extension

Now let's test the extension:

1. Press F5 to launch a new VS Code window with your extension loaded
2. Open a JavaScript/TypeScript file with some functions

3. Configure your OpenAI API key:

   • Open Settings (Cmd/Ctrl + ,)
   • Search for "AI Documenter"
   • Enter your API key

4. Select some code and right-click to see "Generate AI Comment" in the context menu

Example Usage

Let's test with this function:

function calculateCompoundInterest(principal, rate, time, frequency) {
    const amount = principal * Math.pow(1 + (rate / frequency), frequency * time);
    return amount - principal;
}

Select the function and run "Generate AI Comment". The extension should generate something like:

/**
 * Calculates compound interest on a principal amount
 * @param {number} principal - The initial amount of money
 * @param {number} rate - The annual interest rate (as a decimal)
 * @param {number} time - The time period in years
 * @param {number} frequency - How many times interest is compounded per year
 * @returns {number} The compound interest earned
 */
function calculateCompoundInterest(principal, rate, time, frequency) {
    const amount = principal * Math.pow(1 + (rate / frequency), frequency * time);
    return amount - principal;
}

Packaging Your Extension

To package your extension for distribution:

npm install -g vsce
vsce package

This creates a .vsix file you can install locally or publish to the VS Code marketplace.

What's Next?

This extension provides a solid foundation, but you could enhance it further:

• Add support for more programming languages
• Implement smarter function detection using AST parsing
• Add batch documentation for entire files
• Include error handling improvements
• Add support for different documentation formats (Sphinx, Doxygen, etc.)

Conclusion

You've built a functional VS Code extension that demonstrates practical AI integration. The extension shows how AI can genuinely improve developer productivity by handling routine documentation tasks while you focus on solving problems.

The key takeaway isn't just about the specific implementation—it's about understanding how to integrate AI capabilities into familiar tools. This pattern of enhancing existing workflows with AI assistance, rather than replacing them entirely, often provides the most value in real-world development scenarios.

Try the extension with your own code and see how it performs. You mig

---

Tools mentioned:

• Amazon