LLM-Powered VS Code Extension for Context-Aware Code Completion

CodeCom/DEEPCC is a research-driven VS Code extension that enhances traditional IntelliSense with LLM-based, structure-aware code completion.
It integrates program analysis, offset-based inference, and a persistent Python prediction server to deliver intelligent, context-aware suggestions inside Java source files.

This project was developed as part of a broader research effort on LLM-assisted developer tools and is designed to be extensible for tasks such as clone detection, summarization, and refactoring.

Key Features

LLM-based code completion (Java)
Structure-aware predictions using method-level context
Offset-based inference (no visible [MSK] tokens in editor)
Persistent Python prediction server (low latency)
IntelliSense fallback via virtual documents
Tree View UI for ranked candidates
WebView review panel to preview & apply completions
Clone detection support (research extension)
Modular, extensible architecture

High-Level Architecture

VS Code Editor → VS Code Extension (Node.js) → input.json → Python Prediction Server → pred_output.json → Tree View / WebView / Completion Dropdown

Project Structure

CodeCom_VSCode/
├── src/                      # VS Code extension core
│   ├── commands/             # Command handlers
│   ├── providers/            # Language completion providers
│   ├── server/               # (reserved for future server logic)
│   ├── shared/               # Shared state & TreeDataProvider
│   ├── ui/                   # WebView panels
│   └── utils/                # Utilities (file IO, hashing, server mgmt)
│
├── antlr_parsers/             # Java parsing (ANTLR)
├── parser/                    # AST & Tree-sitter parsers
├── models/                    # LLM models
├── evaluator/                 # Evaluation scripts
├── data/                      # Datasets
├── output/                    # Prediction outputs
│
├── predictions_server.py      # Persistent Python inference server
├── predictions.py             # Prediction pipeline
├── model.py                   # Model loading & inference
│
├── extension.js               # Extension entry point
├── package.json               # Extension metadata
└── README.md

How Code Completion Works

CodeCom/DEEPCC's code completion pipeline is designed to provide context-aware, LLM-powered suggestions while preserving the responsiveness and reliability of traditional IntelliSense.

1. Triggering a Prediction

Code completion can be triggered in multiple ways:

Using the keyboard shortcut Ctrl + Shift + Space
From the right-click context menu and select "Code Completion by AI Assistant"

When triggered, CodeCom/DEEPCC identifies the cursor position or selected region and initiates the completion workflow.

2. Context Extraction and Offset Identification

Once triggered:

The enclosing Java method is extracted as the prediction context
The cursor offset within the method is recorded
The method-level code and offset information are serialized into an input.json file

This offset-based design allows CodeCom/DEEPCC to perform inference without modifying the editor text or inserting visible mask tokens.

3. LLM-Based Inference (Python Backend)

The serialized context is sent to a persistent Python prediction server, which:

Loads the language model once at startup to avoid cold-start overhead
Performs offset-based or mask-based inference, depending on configuration
Generates the top-k candidate completions
Assigns a confidence score to each candidate
Writes the ranked results to pred_output.json

4. Candidate Ranking and UI Presentation

The VS Code extension reads the prediction results and displays them through multiple UI components:

Completion dropdown at the cursor position
Tree View panel for browsing candidates across multiple prediction points
WebView review panel for previewing and selectively applying completions

Candidates are ranked by model confidence and presented ahead of IntelliSense suggestions.

5. IntelliSense Fallback via Virtual Documents

To ensure robustness and usability, CodeCom/DEEPCC integrates native IntelliSense as a fallback:

A virtual Java document is created in memory using a custom URI scheme
A temporary placeholder token is inserted only in the virtual document
VS Code’s native IntelliSense runs on this virtual representation
IntelliSense results are merged after LLM-generated suggestions

This approach preserves a clean user experience while combining LLM intelligence with traditional IDE support.

How Code Clone Detection Works

In addition to code completion, CodeCom/DEEPCC supports semantic code clone detection, enabling developers to identify logically similar code fragments even when they differ syntactically. This feature is powered by GraphCodeBERT and integrates seamlessly into the VS Code workflow.

Given a selected Java method, CodeCom/DEEPCC treats it as a query and searches for semantically similar code snippets within a user-specified path (e.g., project directory or dataset). Instead of relying on exact text matching, the system compares code based on semantic meaning and structural behavior.

At a high level, the clone detection workflow is as follows:

The developer selects a Java method in the editor and triggers Code Clone detection.
CodeCom/DEEPCC analyzes the selected method and captures its semantic representation.
The system searches the specified path for other code snippets with similar semantic characteristics.
Matching snippets are ranked by similarity and presented to the user in a WebView panel within VS Code.

The retrieved code clones may:

Use different variable names
Follow different syntactic structures
Appear stylistically different

Yet, they achieve the same underlying functionality as the query method.

This capability helps developers:

Discover related implementations
Reuse existing logic
Understand alternative solutions
Improve consistency across a codebase

By focusing on semantic similarity rather than surface-level syntax, CodeCom/DEEPCC enables more accurate and meaningful clone detection than traditional text-based approaches.

6. Developer Benefits

Code clone detection enables developers to:

Quickly locate related implementations
Reuse proven logic
Understand alternative coding patterns
Improve consistency across a codebase
Accelerate onboarding and code comprehension

By combining graph-based program representations with LLM-driven semantic embeddings, CodeCom/DEEPCC extends beyond code completion to provide deep program understanding capabilities.
This feature demonstrates the extensibility of CODECOM’s architecture for advanced developer assistance tasks such as clone detection, refactoring, and code search.

Video Tutorial

A step-by-step video tutorial demonstrating CODECOM’s features, including code completion and code clone detection, is available here:

▶ Video Tutorial:
Running plugin

The video covers:

Running the VS Code extension
Triggering LLM-based code completion
Using Tree View and WebView panels
Running semantic code clone detection
Applying predictions inside the editor

Model Pretraining, Parsing, and Fine-Tuning

It is powered by a Transformer-based large language model (CodeLM) designed to understand source code at syntactic, semantic, and structural levels.

Pretraining with Program-Aware Tasks

Unlike traditional language models trained only on token sequences, the model is pretrained using multiple code-specific objectives that explicitly incorporate program structure and dependencies.

The model is pretrained using a sequence-to-sequence (Seq2Seq) Transformer architecture, where both the encoder and decoder jointly learn from code tokens, Abstract Syntax Trees (ASTs), and Program Dependence Graphs (PDGs).

Pretraining Tasks

We introduced four specialized pretraining tasks:

1. Masked Sequence-to-Sequence (MSS)

Random code spans are masked
The model reconstructs missing code fragments
Enhances syntactic understanding and long-range context modeling

2. AST Node Prediction (ANP)

Code snippets are paired with either correct or incorrect ASTs
The model learns to distinguish valid syntactic structures
Improves awareness of program grammar and structure

3. Program Dependence Path Prediction (PDPP)

Dependency paths in Program Dependence Graphs are masked
The model reconstructs missing dependency relationships
Enables learning of data flow and control flow semantics

4. Code Functionality Prediction (CFP)

Functional elements (e.g., method names, API calls) are masked
The model predicts missing functionality based on context
Strengthens semantic intent understanding

Role of Parsing in Learning

What Parsing Means

Parsing is the process of transforming raw source code into structured representations such as:

Abstract Syntax Trees (ASTs) – representing syntactic structure
Control Flow Graphs (CFGs) – representing execution flow
Program Dependence Graphs (PDGs) – representing data and control dependencies

These representations allow the model to reason about how code behaves, not just how it looks.

How Parsing Is Used in model

Parsing plays a central role throughout the model pipeline. Parsing is the process of analyzing source code to understand its syntactic structure and represent it in a structured form such as an Abstract Syntax Tree (AST). This structured representation enables tools to reason about:

AST parsing (via Tree-sitter / ANTLR) extracts hierarchical structure
Program slicing and data-flow analysis identify def–use relationships
Program dependence analysis captures semantic interactions between code elements

These parsed structures are serialized and embedded into the model input during pretraining, enabling model to learn structure-aware representations instead of plain token sequences.

Parsing Technologies Used

ANTLR – Grammar-based parsing for Java
Tree-sitter – Incremental parsing and AST generation

Fine-Tuning for Code Completion

After pretraining, the model is fine-tuned specifically for context-aware code completion.

Fine-Tuning Data

Large-scale Java and Python datasets
Each sample pairs incomplete code with the correct completion
Prediction points are identified using offset-based analysis

Fine-Tuning Objective

Given a method-level context and prediction offset:
- Generate top-k candidate completions
- Rank candidates by confidence
Optimized using cross-entropy loss with evaluation on BLEU and Accuracy@k

Fine-tuning aligns the pretrained representations with real-world IDE usage, enabling the model to generate accurate completions without inserting visible mask tokens into the editor.

Model Deployment and Availability

The final fine-tuned model used in CodeCom/DEEPCC is publicly available on our Hugging Face, ensuring transparency and reproducibility.

Hugging Face Model:
Model Link

Users can:

Run inference independently
Fine-tune the model for other languages or tasks
Integrate the model into custom developer tools

Reproducing Results End-to-End

To reproduce the model's predictions:

Download the model and vocabularies from Hugging Face
Use the provided preprocessing and parsing pipeline
Run inference with the fixed configuration
Compare ranked predictions and confidence scores

Because all components like model, vocabularies, configuration, and decoding strategy are versioned and fixed, CodeCom/DEEPCC supports transparent and repeatable experimentation.

Installation & Setup (Prerequisites)

VS Code (latest)
Node.js ≥ 18
Python ≥ 3.9
PyTorch
Hugging Face Transformers

Clone the repository git clone cd CodeCom_VSCode
Install Extension Dependencies npm install
Install Python Dependencies pip install -r requirements.txt
Start Prediction Server python predictions_server.py

The server is also auto-launched when the extension activates.

Run Extension in VS Code npm run compile

Press F5 to launch the Extension Development Host.

Getting Started with VS Code Extension Development

For users who are new to VS Code extension development or would like to understand the fundamentals, the official VS Code documentation provides an excellent starting point.

Official VS Code Extension Guide:
Getting started with VSCode

This tutorial covers:

Setting up the VS Code extension development environment
Understanding the extension lifecycle (activate, deactivate)
Registering commands and contributions
Running and debugging extensions locally

Usage Guide (Code Completion)

Open a Java file
Place cursor inside a method
Trigger completion by pressing Ctrl + Shift + Space or right-click → Code Completion by AI Assistant
Review suggestions in dropdown or Tree View
Preview in WebView
Apply selected completion

Usage Guide (Clone Detection)

Open a Java file
Select a method or a class
Right-click → Clone Detection by AI Assistant
Review suggestions in WebView

Datasets Used

This project relies on large-scale, publicly available code datasets to support model pretraining, fine-tuning, and evaluation.
All datasets are derived from open-source repositories and are processed to preserve syntactic correctness, structural information, and semantic integrity.

Pretraining Datasets

For pretraining the DeepCC model, large-scale source code corpora were used to expose the model to diverse programming patterns, APIs, and coding styles.

Source

CodeSearchNet Dataset
- Languages: Java, Python (and others)
- Contains millions of functions extracted from open-source GitHub repositories
- Includes paired code and documentation

Dataset Link:
Datasets Used

All datasets are stored under the data/datasets/ directory and are organized by language and training phase.

data/datasets/
├── Java_pretrain_dataset/
├── Java_finetune_dataset/
├── Python_pretrain_dataset/
└── Python_finetune_dataset/

During preprocessing:

Source files were parsed into Abstract Syntax Trees (ASTs)
Program dependence information (control and data flow) was extracted
Code samples failing to parse correctly were filtered out

This ensured high-quality, structurally valid inputs for multi-task pretraining.

Fine-Tuning Datasets

For task-specific fine-tuning, curated subsets of Java and Python code were used to align the pretrained model with IDE-style code completion scenarios.

Fine-Tuning Characteristics

Method-level code snippets
Created prediction points using offset-based masking
Ground-truth completions preserved for supervised learning

These datasets emphasize realistic developer workflows, where code is partially written and requires intelligent continuation.

Evaluation Datasets

Model evaluation was performed using held-out code samples that were not seen during training.

Evaluation datasets were used to compute:

BLEU score
Accuracy@k
Top-k prediction success rates

This separation ensures fair evaluation and avoids data leakage.

Dataset Availability and Reproducibility

All dataset sources used in this project are publicly available.

Where redistribution is permitted, processed dataset samples and scripts for data preparation are included in this repository.
Where redistribution is restricted, scripts and instructions are provided to regenerate the datasets from the original sources.

Dataset Resources

CodeSearchNet:
Java Dataset
Py150: Python Dataset

Notes on Reproducibility

To reproduce the datasets used in this project:

Download the original dataset from the provided links
Run the preprocessing scripts included in this repository
Parse code samples using the specified parsing tools
Generate training and evaluation splits as described in the documentation

This approach ensures transparent, repeatable, and extensible experimentation while respecting dataset licensing constraints.

Updating Dataset Paths

To reproduce training or inference, users only need to update the dataset paths in args.py (CODE-Completion-LLM/code-completion/sources/args.py).

Example configuration:

pretrain_file_path: str = field( default="/../../data/datasets/Python_pretrain_dataset", metadata={"help": "Path to pretraining files"} )

finetune_file_path: str = field( default="/../../data/datasets/Python_finetune_dataset", metadata={"help": "Path to finetuning files"} )

No code changes are required beyond updating these paths.

Research Paper & Citation

This project is part of a published research effort.

Paper:
Intelligent Code Completion by a Unified Multi-task Learning with a Large Language Model
2025 IEEE/ACIS 23rd International Conference on Software Engineering Research, Management and Applications (SERA)

Citation:

@INPROCEEDINGS{11154531,
  author={Maharjan, Shradha and Xia, Meng and Ahn, Tae Hyuk and Song, Myoungkyu},
  booktitle={2025 IEEE/ACIS 23rd International Conference on Software Engineering Research, Management and Applications (SERA)}, 
  title={Intelligent Code Completion by a Unified Multi-task Learning with a Large Language Model}, 
  year={2025},
  volume={},
  number={},
  pages={134-142},
  keywords={Measurement;Codes;Accuracy;Large language models;Source coding;Software quality;Syntactics;Multitasking;Software development management;Software engineering},
  doi={10.1109/SERA65747.2025.11154531}}

Name		Name	Last commit message	Last commit date
Latest commit History 3 Commits
antlr_parsers		antlr_parsers
data		data
evaluator		evaluator
models		models
node_modules		node_modules
parser		parser
src		src
test		test
.DS_Store		.DS_Store
.vscode-test.mjs		.vscode-test.mjs
.vscodeignore		.vscodeignore
CHANGELOG.md		CHANGELOG.md
README.md		README.md
enums.py		enums.py
eslint.config.mjs		eslint.config.mjs
extension.js		extension.js
git-proc.sh		git-proc.sh
java_class_method_visitor.py		java_class_method_visitor.py
java_treesitter_parser.py		java_treesitter_parser.py
model.py		model.py
predictions.py		predictions.py
predictions_server.py		predictions_server.py
registerCompletion.js		registerCompletion.js
run.py		run.py
run_v2.py		run_v2.py
util_ast.py		util_ast.py
vsc-extension-quickstart.md		vsc-extension-quickstart.md

Uh oh!

Folders and files

Latest commit

History

Repository files navigation