fix(ts-plugin): find Astro.locals references

.changeset/fix-ts-plugin-astro-locals-references.md

@@ -0,0 +1,5 @@
+---
+'@astrojs/ts-plugin': patch
+---
+
+Fixes Go To References for members accessed through `Astro.locals` in Astro files.

packages/language-tools/ts-plugin/.vscode-test.mjs

@@ -5,6 +5,7 @@ export default defineConfig([
 		label: 'unitTests',
 		files: 'test/**/*.test.mts',
 		extensionDevelopmentPath: '../vscode',
+		workspaceFolder: './test/fixtures',
 		version: 'stable',
 		mocha: {
 			ui: 'tdd',

packages/language-tools/ts-plugin/src/astro-types.ts

@@ -0,0 +1,126 @@
+import path from 'node:path';
+import type ts from 'typescript';
+
+const decoratedHosts = new WeakSet<ts.LanguageServiceHost>();
+const decoratedProjects = new WeakSet<ts.LanguageServiceHost>();
+
+export function findAstroPackageDirectory(
+	tsModule: typeof import('typescript'),
+	currentDirectory: string | string[],
+): string | undefined {
+	for (const candidate of Array.isArray(currentDirectory) ? currentDirectory : [currentDirectory]) {
+		const astroDirectory = findAstroPackageDirectoryFrom(tsModule, candidate);
+		if (astroDirectory) {
+			return astroDirectory;
+		}
+	}
+}
+
+function findAstroPackageDirectoryFrom(
+	tsModule: typeof import('typescript'),
+	currentDirectory: string,
+): string | undefined {
+	let directory = tsModule.sys.resolvePath(currentDirectory);
+
+	while (true) {
+		const packageJson = path.join(directory, 'node_modules', 'astro', 'package.json');
+		if (tsModule.sys.fileExists(packageJson)) {
+			return path.dirname(packageJson);
+		}
+
+		const parent = path.dirname(directory);
+		if (parent === directory) {
+			return undefined;
+		}
+		directory = parent;
+	}
+}
+
+export function addAstroTypes(
+	tsModule: typeof import('typescript'),
+	host: ts.LanguageServiceHost,
+	currentDirectory: string | string[],
+) {
+	if (decoratedHosts.has(host)) {
+		return;
+	}
+
+	const astroDirectory = findAstroPackageDirectory(tsModule, currentDirectory);
+	if (!astroDirectory) {
+		return;
+	}
+
+	const addedFileNames = ['./env.d.ts', './astro-jsx.d.ts']
+		.map((filePath) => tsModule.sys.resolvePath(path.resolve(astroDirectory, filePath)))
+		.filter((fileName) => tsModule.sys.fileExists(fileName));
+
+	if (!addedFileNames.length) {
+		return;
+	}
+
+	decoratedHosts.add(host);
+
+	const getScriptFileNames = host.getScriptFileNames.bind(host);
+	host.getScriptFileNames = () => {
+		const fileNames = getScriptFileNames();
+		const seen = new Set(fileNames);
+
+		return [...fileNames, ...addedFileNames.filter((fileName) => !seen.has(fileName))];
+	};
+}
+
+export function addAstroProjectFiles(
+	tsModule: typeof import('typescript'),
+	project: ts.server.Project,
+	host: ts.LanguageServiceHost,
+) {
+	if (
+		decoratedProjects.has(host) ||
+		project.projectKind !== tsModule.server.ProjectKind.Configured
+	) {
+		return;
+	}
+	decoratedProjects.add(host);
+
+	const getScriptFileNames = host.getScriptFileNames.bind(host);
+	let lastProjectVersion: string | undefined;
+	let astroFiles: string[] = [];
+
+	host.getScriptFileNames = () => {
+		const fileNames = getScriptFileNames();
+		const projectVersion = host.getProjectVersion?.();
+		if (!astroFiles.length || projectVersion !== lastProjectVersion) {
+			lastProjectVersion = projectVersion;
+			astroFiles = findAstroFiles(tsModule, project);
+		}
+
+		const seen = new Set(fileNames);
+		return [...fileNames, ...astroFiles.filter((fileName) => !seen.has(fileName))];
+	};
+}
+
+function findAstroFiles(
+	tsModule: typeof import('typescript'),
+	project: ts.server.Project,
+): string[] {
+	const configFile = project.getProjectName();
+	const config = tsModule.readJsonConfigFile(configFile, project.readFile.bind(project));
+	const parseHost = {
+		useCaseSensitiveFileNames: project.useCaseSensitiveFileNames(),
+		fileExists: project.fileExists.bind(project),
+		readFile: project.readFile.bind(project),
+		readDirectory: (...args: Parameters<ts.server.Project['readDirectory']>) => {
+			args[1] = ['.astro'];
+			return project.readDirectory(...args);
+		},
+	};
+	const parsed = tsModule.parseJsonSourceFileConfigFileContent(
+		config,
+		parseHost,
+		project.getCurrentDirectory(),
+		undefined,
+		configFile,
+	);
+
+	return parsed.fileNames;
+}

packages/language-tools/ts-plugin/src/index.ts

@@ -1,14 +1,22 @@
 import type { LanguagePlugin } from '@volar/language-core';
 import { createLanguageServicePlugin } from '@volar/typescript/lib/quickstart/createLanguageServicePlugin.js';
+import path from 'node:path';
+import { addAstroProjectFiles, addAstroTypes } from './astro-types.js';
 import type { CollectionConfig } from './frontmatter.js';
 import { getFrontmatterLanguagePlugin } from './frontmatter.js';
 import { getLanguagePlugin } from './language.js';
 
 export = createLanguageServicePlugin((ts, info) => {
 	let collectionConfig = undefined;
+	const currentDir = info.project.getCurrentDirectory();
+
+	addAstroProjectFiles(ts, info.project, info.languageServiceHost);
+	addAstroTypes(ts, info.languageServiceHost, [
+		currentDir,
+		...info.languageServiceHost.getScriptFileNames().map((fileName) => path.dirname(fileName)),
+	]);
 
 	try {
-		const currentDir = info.project.getCurrentDirectory();
 		const fileContent = ts.sys.readFile(currentDir + '/.astro/collections/collections.json');
 		if (fileContent) {
 			collectionConfig = {

packages/language-tools/ts-plugin/test/fixtures/LocalsReference.astro

@@ -0,0 +1 @@
+<div>{Astro.locals.utils.toUpper("Astro")}</div>

packages/language-tools/ts-plugin/test/fixtures/astro.config.mjs

@@ -0,0 +1 @@
+export default {};

packages/language-tools/ts-plugin/test/fixtures/locals.d.ts

@@ -0,0 +1,9 @@
+export {};
+
+declare global {
+	namespace App {
+		interface Locals {
+			utils: import('./utilClass').Utils;
+		}
+	}
+}

packages/language-tools/ts-plugin/test/fixtures/package.json

@@ -0,0 +1,6 @@
+{
+	"private": true,
+	"devDependencies": {
+		"astro": "link:./test-astro"
+	}
+}

packages/language-tools/ts-plugin/test/fixtures/test-astro/astro-jsx.d.ts

@@ -0,0 +1,3 @@
+declare namespace astroHTML.JSX {
+	interface HTMLAttributes {}
+}

packages/language-tools/ts-plugin/test/fixtures/test-astro/env.d.ts

@@ -0,0 +1,6 @@
+type AstroGlobal = {
+	locals: App.Locals;
+};
+
+declare const Astro: Readonly<AstroGlobal>;
+declare const Fragment: any;

packages/language-tools/ts-plugin/test/fixtures/test-astro/package.json

@@ -0,0 +1,5 @@
+{
+	"name": "test-astro-fixture",
+	"private": true,
+	"version": "0.0.0"
+}

packages/language-tools/ts-plugin/test/fixtures/utilClass.ts

@@ -0,0 +1,5 @@
+export class Utils {
+	toUpper(value: string) {
+		return value.toUpperCase();
+	}
+}

packages/language-tools/ts-plugin/test/suite/extension.test.mts

@@ -39,6 +39,23 @@ suite('Extension Test Suite', () => {
 		assert.strictEqual(hasAstroRef, true, 'Should find Astro reference');
 	}).timeout(50000);
 
+	test('can find references through Astro.locals inside Astro files', async () => {
+		const doc = await vscode.workspace.openTextDocument(
+			vscode.Uri.file(path.join(__dirname, '../fixtures/utilClass.ts')),
+		);
+
+		const references = await waitForTS<vscode.Location[]>(
+			'vscode.executeReferenceProvider',
+			[doc.uri, doc.positionAt(doc.getText().indexOf('toUpper'))],
+			(result) => result.some((ref) => ref.uri.path.includes('LocalsReference.astro')),
+		);
+
+		const hasAstroLocalsRef = references.some((ref) =>
+			ref.uri.path.includes('LocalsReference.astro'),
+		);
+		assert.strictEqual(hasAstroLocalsRef, true, 'Should find Astro.locals reference');
+	}).timeout(50000);
+
 	test('can get completions for Astro components', async () => {
 		const doc = await vscode.workspace.openTextDocument(
 			vscode.Uri.file(path.join(__dirname, '../fixtures/script.ts')),

packages/language-tools/ts-plugin/test/units/astro-types.test.mts

@@ -0,0 +1,129 @@
+import 'mocha';
+import assert from 'node:assert';
+import fs from 'node:fs';
+import os from 'node:os';
+import path from 'node:path';
+import ts from 'typescript';
+import { addAstroTypes } from '../../src/astro-types.js';
+import { astro2tsx } from '../../src/astro2tsx.js';
+
+function createFixture() {
+	const root = fs.mkdtempSync(path.join(os.tmpdir(), 'astro-ts-plugin-'));
+	const src = path.join(root, 'src');
+	const astroPackage = path.join(root, 'node_modules', 'astro');
+
+	fs.mkdirSync(src, { recursive: true });
+	fs.mkdirSync(astroPackage, { recursive: true });
+
+	const utilsFile = path.join(src, 'utils.ts');
+	const envFile = path.join(src, 'env.d.ts');
+	const astroFile = path.join(src, 'index.astro');
+	const astroTsxFile = path.join(src, 'index.astro.tsx');
+
+	fs.writeFileSync(path.join(astroPackage, 'package.json'), '{"name":"astro","version":"6.0.0"}');
+	fs.writeFileSync(
+		path.join(astroPackage, 'env.d.ts'),
+		[
+			'type AstroGlobal = { locals: App.Locals };',
+			'declare const Astro: Readonly<AstroGlobal>;',
+			'declare const Fragment: any;',
+		].join('\n'),
+	);
+	fs.writeFileSync(path.join(astroPackage, 'astro-jsx.d.ts'), '');
+	fs.writeFileSync(path.join(astroPackage, 'jsx-runtime.d.ts'), 'export {};');
+	fs.writeFileSync(
+		utilsFile,
+		[
+			'export class Utils {',
+			'\ttoUpper(value: string) {',
+			'\t\treturn value.toUpperCase();',
+			'\t}',
+			'}',
+		].join('\n'),
+	);
+	fs.writeFileSync(
+		envFile,
+		[
+			'export {};',
+			'declare global {',
+			'\tnamespace App {',
+			'\t\tinterface Locals {',
+			'\t\t\tutils: import("./utils").Utils;',
+			'\t\t}',
+			'\t}',
+			'}',
+		].join('\n'),
+	);
+	fs.writeFileSync(astroFile, '<div>{Astro.locals.utils.toUpper("Astro")}</div>\n');
+	const astroTsx = astro2tsx(fs.readFileSync(astroFile, 'utf8'), astroFile);
+	fs.writeFileSync(
+		astroTsxFile,
+		astroTsx.virtualFile.snapshot.getText(0, astroTsx.virtualFile.snapshot.getLength()),
+	);
+
+	return { root, files: [utilsFile, envFile, astroTsxFile], utilsFile, astroTsxFile };
+}
+
+function findToUpperReferenceFiles(injectAstroTypes: boolean) {
+	const fixture = createFixture();
+	const versions = new Map(fixture.files.map((fileName) => [fileName, '0']));
+	const host: ts.LanguageServiceHost = {
+		getScriptFileNames: () => fixture.files,
+		getScriptVersion: (fileName) => versions.get(fileName) ?? '0',
+		getScriptSnapshot: (fileName) => {
+			if (!fs.existsSync(fileName)) {
+				return undefined;
+			}
+			return ts.ScriptSnapshot.fromString(fs.readFileSync(fileName, 'utf8'));
+		},
+		getCurrentDirectory: () => fixture.root,
+		getCompilationSettings: () => ({
+			allowJs: true,
+			jsx: ts.JsxEmit.Preserve,
+			module: ts.ModuleKind.ESNext,
+			moduleResolution: ts.ModuleResolutionKind.Node10,
+			target: ts.ScriptTarget.ESNext,
+		}),
+		getDefaultLibFileName: (options) => ts.getDefaultLibFilePath(options),
+		fileExists: ts.sys.fileExists,
+		readFile: ts.sys.readFile,
+		readDirectory: ts.sys.readDirectory,
+		directoryExists: ts.sys.directoryExists,
+		getDirectories: ts.sys.getDirectories,
+		useCaseSensitiveFileNames: () => ts.sys.useCaseSensitiveFileNames,
+		getNewLine: () => ts.sys.newLine,
+		getScriptKind: (fileName) => (fileName.endsWith('.tsx') ? ts.ScriptKind.TSX : ts.ScriptKind.TS),
+	};
+
+	if (injectAstroTypes) {
+		addAstroTypes(ts, host, fixture.root);
+	}
+
+	const service = ts.createLanguageService(host);
+	const utilsText = fs.readFileSync(fixture.utilsFile, 'utf8');
+	const references = service.findReferences(fixture.utilsFile, utilsText.indexOf('toUpper')) ?? [];
+
+	try {
+		return references.flatMap((entry) =>
+			entry.references.map((reference) => path.relative(fixture.root, reference.fileName)),
+		);
+	} finally {
+		fs.rmSync(fixture.root, { recursive: true, force: true });
+	}
+}
+
+suite('Astro type injection', () => {
+	test('adds Astro globals so references through Astro.locals are discoverable', () => {
+		const referencesWithoutAstroTypes = findToUpperReferenceFiles(false);
+		assert.ok(
+			!referencesWithoutAstroTypes.includes(path.join('src', 'index.astro.tsx')),
+			'fixture should reproduce missing Astro.locals references before injecting Astro types',
+		);
+
+		const referencesWithAstroTypes = findToUpperReferenceFiles(true);
+		assert.ok(
+			referencesWithAstroTypes.includes(path.join('src', 'index.astro.tsx')),
+			'should find the Astro.locals reference in the Astro file',
+		);
+	});
+});