From eb49fb267d2615a4023a444ffbdbaab75fb9ccf2 Mon Sep 17 00:00:00 2001
From: Erika <3019731+Princesseuh@users.noreply.github.com>
Date: Fri, 22 Mar 2024 18:06:29 +0100
Subject: [PATCH] fix(html): Load custom data the proper way (#841)

* fix(html): Load custom data the proper way

* chore: changeset
---
 .changeset/thirty-peas-turn.md                |  6 ++++
 packages/language-server/src/plugins/html.ts  | 28 ++++++++++++++-----
 .../test/html/custom-data.test.ts             | 22 +++++++++++++++
 3 files changed, 49 insertions(+), 7 deletions(-)
 create mode 100644 .changeset/thirty-peas-turn.md
 create mode 100644 packages/language-server/test/html/custom-data.test.ts

diff --git a/.changeset/thirty-peas-turn.md b/.changeset/thirty-peas-turn.md
new file mode 100644
index 00000000..11cf70b8
--- /dev/null
+++ b/.changeset/thirty-peas-turn.md
@@ -0,0 +1,6 @@
+---
+"@astrojs/language-server": patch
+"astro-vscode": patch
+---
+
+Fixes completions for Astro-specific attributes not working in certain contexts
diff --git a/packages/language-server/src/plugins/html.ts b/packages/language-server/src/plugins/html.ts
index 826f71b7..7d6ea23a 100644
--- a/packages/language-server/src/plugins/html.ts
+++ b/packages/language-server/src/plugins/html.ts
@@ -1,22 +1,36 @@
 import { CompletionItemKind, ServicePlugin, ServicePluginInstance } from '@volar/language-server';
 import { create as createHtmlService } from 'volar-service-html';
+import * as html from 'vscode-html-languageservice';
+import { URI, Utils } from 'vscode-uri';
 import { AstroVirtualCode } from '../core/index.js';
 import { astroAttributes, astroElements, classListAttribute } from './html-data.js';
 import { isInComponentStartTag } from './utils.js';
 
 export const create = (): ServicePlugin => {
-	const htmlServicePlugin = createHtmlService();
+	const htmlServicePlugin = createHtmlService({
+		getCustomData: async (context) => {
+			const customData: string[] = (await context.env.getConfiguration?.('html.customData')) ?? [];
+			const newData: html.IHTMLDataProvider[] = [];
+			for (const customDataPath of customData) {
+				const uri = Utils.resolvePath(URI.parse(context.env.workspaceFolder), customDataPath);
+				const json = await context.env.fs?.readFile?.(uri.toString());
+				if (json) {
+					try {
+						const data = JSON.parse(json);
+						newData.push(html.newHTMLDataProvider(customDataPath, data));
+					} catch (error) {
+						console.error(error);
+					}
+				}
+			}
+			return [...newData, astroAttributes, astroElements, classListAttribute];
+		},
+	});
 	return {
 		...htmlServicePlugin,
 		create(context): ServicePluginInstance {
 			const htmlPlugin = htmlServicePlugin.create(context);
 
-			htmlPlugin.provide['html/updateCustomData']?.([
-				astroAttributes,
-				astroElements,
-				classListAttribute,
-			]);
-
 			return {
 				...htmlPlugin,
 				async provideCompletionItems(document, position, completionContext, token) {
diff --git a/packages/language-server/test/html/custom-data.test.ts b/packages/language-server/test/html/custom-data.test.ts
new file mode 100644
index 00000000..e67a7251
--- /dev/null
+++ b/packages/language-server/test/html/custom-data.test.ts
@@ -0,0 +1,22 @@
+import { Position } from '@volar/language-server';
+import { expect } from 'chai';
+import { describe } from 'mocha';
+import { type LanguageServer, getLanguageServer } from '../server.js';
+
+describe('HTML - Custom Data', () => {
+	let languageServer: LanguageServer;
+
+	before(async () => (languageServer = await getLanguageServer()));
+
+	it('Can properly get completions for attributes added by our custom data', async () => {
+		const document = await languageServer.openFakeDocument(`<div class></div>`, 'astro');
+		const completions = await languageServer.handle.sendCompletionRequest(
+			document.uri,
+			Position.create(0, 10)
+		);
+
+		const labels = completions!.items.map((i) => i.label);
+		expect(completions!.items).to.not.be.empty;
+		expect(labels).to.include('class:list');
+	});
+});