feat(cdk-explorer): add cdk/getConstructTree LSP route

packages/@aws-cdk/cdk-explorer/lib/lsp/construct-tree-request.ts

@@ -0,0 +1,48 @@
+import { RequestType0 } from 'vscode-languageserver';
+import type { SourceLocation } from '../core/source-resolver';
+
+/**
+ * Custom LSP request that returns the source-resolved construct tree, flattened
+ * to one entry per construct keyed by its tree path.
+ *
+ * The server already builds this tree from the cloud assembly (for hover,
+ * go-to-definition and CodeLens), so this request lets a client -- for example
+ * the AWS Toolkit's CDK tree -- overlay source links onto its own construct tree
+ * without re-implementing source resolution.
+ */
+export const GET_CONSTRUCT_TREE_METHOD = 'cdk/getConstructTree';
+
+/** One construct, flattened, with its resolved source location. */
+export interface ConstructSourceEntry {
+  /** The construct's tree path (as in tree.json); the key a consumer joins on. */
+  readonly path: string;
+  readonly id: string;
+  /** CFN resource type, if this construct is a CloudFormation resource. */
+  readonly type?: string;
+  /** CFN logical id, if this construct maps to a resource. */
+  readonly logicalId?: string;
+  /** User source location where the construct was created; undefined when unresolved. */
+  readonly sourceLocation?: SourceLocation;
+  /** Absolute path to the construct's template, containment-checked; undefined when none. */
+  readonly templateFile?: string;
+  /**
+   * Character offset of the resource's block within `templateFile`, so a client
+   * can open the template positioned on the resource. Undefined when there is no
+   * template or the resource block can't be located.
+   */
+  readonly templateOffset?: number;
+}
+
+export interface GetConstructTreeResult {
+  /** 'no-assembly' when the app has not been synthesized yet (no cdk.out). */
+  readonly status: 'ok' | 'no-assembly';
+  /** The cdk.out directory the tree was read from. */
+  readonly assemblyDir: string;
+  /** One entry per construct, in pre-order. Empty when status is 'no-assembly'. */
+  readonly entries: readonly ConstructSourceEntry[];
+  /** Non-fatal warnings gathered while resolving (e.g. an unparseable source map). */
+  readonly warnings: readonly string[];
+}
+
+/** No request params: the server serves its single initialized project. */
+export const GetConstructTreeRequest = new RequestType0<GetConstructTreeResult, void>(GET_CONSTRUCT_TREE_METHOD);

packages/@aws-cdk/cdk-explorer/lib/lsp/server.ts

@@ -1,7 +1,7 @@
 import * as fs from 'fs';
 import * as path from 'path';
 import { fileURLToPath } from 'url';
-import { ConstructIndex } from '@aws-cdk/cloud-assembly-api';
+import { ConstructIndex, indexTemplateRanges, type TemplateRanges } from '@aws-cdk/cloud-assembly-api';
 import { ToolkitError } from '@aws-cdk/toolkit-lib';
 import {
   StreamMessageReader,
@@ -28,6 +28,7 @@ import {
 import { WATCH_EXCLUDE_DEFAULTS, createIgnoreMatcher } from '../api-private';
 import { codeLensesForFile } from './codelens';
 import { executeCommand, SUPPORTED_COMMANDS, type NotifySink } from './commands';
+import { GetConstructTreeRequest, type ConstructSourceEntry, type GetConstructTreeResult } from './construct-tree-request';
 import { mapViolationsToDiagnostics } from './diagnostics';
 import { hoverForPosition } from './hover';
 import { offsetAtPosition } from './positions';
@@ -137,6 +138,8 @@ export interface LspHandlers {
   onDefinition(params: DefinitionParams): Promise<Location | undefined>;
   onExecuteCommand(params: ExecuteCommandParams): Promise<void>;
   onHover(params: HoverParams): Promise<Hover | undefined>;
+  /** Serve the source-resolved construct tree (flattened) for a client overlay. */
+  onGetConstructTree(): GetConstructTreeResult;
   onShutdown(): void;
 }
 
@@ -208,6 +211,12 @@ export function createLspHandlers(options: LspHandlerOptions): LspHandlers {
   // Latest index from readAssembly, served to CodeLens. Refreshed at startup
   // and whenever the cdk.out watcher detects a re-synth.
   let cachedIndex: ConstructIndex<ConstructNode> = ConstructIndex.fromTree<ConstructNode>([]);
+  // The decorated tree, its warnings, and the source dir behind cachedIndex,
+  // retained so the getConstructTree request serves source locations without a
+  // re-read. undefined assemblyDir means no successful read has happened yet.
+  let cachedTree: readonly ConstructNode[] = [];
+  let cachedWarnings: readonly string[] = [];
+  let cachedAssemblyDir: string | undefined;
   // URIs that currently have published diagnostics. On each refresh we publish
   // an empty array for any URI that no longer has diagnostics, otherwise a
   // resolved violation would leave a stale squiggle behind.
@@ -272,6 +281,9 @@ export function createLspHandlers(options: LspHandlerOptions): LspHandlers {
     }
 
     cachedIndex = ConstructIndex.fromTree(tree);
+    cachedTree = tree;
+    cachedWarnings = warnings;
+    cachedAssemblyDir = assemblyDir;
 
     const { byUri, dropped } = mapViolationsToDiagnostics(violations, cachedIndex);
 
@@ -479,6 +491,46 @@ export function createLspHandlers(options: LspHandlerOptions): LspHandlers {
       return hoverForPosition(cachedIndex, params.textDocument.uri, params.position,
         (file) => fs.promises.readFile(file, 'utf-8').catch(() => undefined));
     },
+    onGetConstructTree() {
+      const assemblyDir = cachedAssemblyDir ?? path.join(currentProjectDir(), 'cdk.out');
+      if (cachedAssemblyDir === undefined) {
+        return { status: 'no-assembly', assemblyDir, entries: [], warnings: [] };
+      }
+      // Flatten the already-decorated tree in pre-order; the client joins on `path`.
+      // Parse each template's resource ranges at most once per request.
+      const rangesByTemplate = new Map<string, TemplateRanges | undefined>();
+      const templateRangesFor = (templateFile: string): TemplateRanges | undefined => {
+        if (!rangesByTemplate.has(templateFile)) {
+          try {
+            rangesByTemplate.set(templateFile, indexTemplateRanges(fs.readFileSync(templateFile, 'utf-8')));
+          } catch {
+            rangesByTemplate.set(templateFile, undefined);
+          }
+        }
+        return rangesByTemplate.get(templateFile);
+      };
+      const entries: ConstructSourceEntry[] = [];
+      const collect = (nodes: readonly ConstructNode[]): void => {
+        for (const node of nodes) {
+          const templateOffset =
+            node.templateFile && node.logicalId
+              ? templateRangesFor(node.templateFile)?.block(node.logicalId)?.start
+              : undefined;
+          entries.push({
+            path: node.path,
+            id: node.id,
+            type: node.type,
+            logicalId: node.logicalId,
+            sourceLocation: node.sourceLocation,
+            templateFile: node.templateFile,
+            templateOffset,
+          });
+          collect(node.children);
+        }
+      };
+      collect(cachedTree);
+      return { status: 'ok', assemblyDir, entries, warnings: [...cachedWarnings] };
+    },
     onShutdown() {
       shutdownRequested = true;
       void assemblyWatcher?.close();
@@ -554,6 +606,7 @@ export function startServer(options: LspServerOptions): void {
   connection.onDefinition((params) => handlers.onDefinition(params));
   connection.onExecuteCommand((params) => handlers.onExecuteCommand(params));
   connection.onHover((params) => handlers.onHover(params));
+  connection.onRequest(GetConstructTreeRequest, () => handlers.onGetConstructTree());
   connection.onShutdown(() => handlers.onShutdown());
   connection.onExit(() => process.exit(0));
 

packages/@aws-cdk/cdk-explorer/test/lsp/server.test.ts

@@ -904,3 +904,98 @@ describe('LSP Server -- synth-failure diagnostics', () => {
     expect(entriesForFile[entriesForFile.length - 1].diagnostics).toEqual([]);
   });
 });
+
+describe('getConstructTree request', () => {
+  const emptyViolations = { version: '1.0.0', pluginReports: [] };
+
+  it('flattens the source-resolved tree in pre-order', async () => {
+    const tree = [
+      {
+        path: 'Stack1',
+        id: 'Stack1',
+        children: [
+          {
+            path: 'Stack1/MyBucket',
+            id: 'MyBucket',
+            type: 'AWS::S3::Bucket',
+            logicalId: 'MyBucket123',
+            templateFile: path.join('/p', 'cdk.out', 'Stack1.template.json'),
+            sourceLocation: { file: '/p/lib/stack.ts', line: 12, column: 5 },
+            children: [],
+          },
+        ],
+      },
+    ];
+    const client = createTestClient({
+      readAssembly: async () => ({ status: 'success', data: { warnings: ['bad source map'], tree, violations: emptyViolations } }),
+    });
+    await initializeClient(client, { applicationDir: '/p' });
+
+    const result = client.handlers.onGetConstructTree();
+
+    expect(result.status).toBe('ok');
+    expect(result.assemblyDir).toBe(path.join('/p', 'cdk.out'));
+    expect(result.warnings).toEqual(['bad source map']);
+    // Pre-order: parent before child.
+    expect(result.entries.map((e) => e.path)).toEqual(['Stack1', 'Stack1/MyBucket']);
+    const bucket = result.entries.find((e) => e.path === 'Stack1/MyBucket');
+    expect(bucket).toMatchObject({
+      id: 'MyBucket',
+      type: 'AWS::S3::Bucket',
+      logicalId: 'MyBucket123',
+      sourceLocation: { file: '/p/lib/stack.ts', line: 12, column: 5 },
+      templateFile: path.join('/p', 'cdk.out', 'Stack1.template.json'),
+    });
+  });
+
+  it('includes the template offset of the resource block', async () => {
+    const dir = fs.mkdtempSync(path.join(os.tmpdir(), 'server-tree-'));
+    const outDir = path.join(dir, 'cdk.out');
+    fs.mkdirSync(outDir, { recursive: true });
+    const templateFile = path.join(outDir, 'Stack1.template.json');
+    const text = JSON.stringify({ Resources: { MyBucket123: { Type: 'AWS::S3::Bucket' } } }, undefined, 1);
+    fs.writeFileSync(templateFile, text);
+    try {
+      const client = createTestClient({
+        readAssembly: async () => ({
+          status: 'success',
+          data: {
+            tree: [{
+              path: 'Stack1/MyBucket',
+              id: 'MyBucket',
+              type: 'AWS::S3::Bucket',
+              logicalId: 'MyBucket123',
+              templateFile,
+              children: [],
+            }],
+            violations: emptyViolations,
+            warnings: [],
+          },
+        }),
+      });
+      await initializeClient(client, { applicationDir: dir });
+
+      const result = client.handlers.onGetConstructTree();
+
+      // The offset is the start of the resource's value block: the first `{`
+      // after the logical id key.
+      const expectedStart = text.indexOf('{', text.indexOf('"MyBucket123"'));
+      const bucket = result.entries.find((e) => e.path === 'Stack1/MyBucket');
+      expect(bucket?.templateOffset).toBe(expectedStart);
+      expect(text[bucket!.templateOffset!]).toBe('{');
+    } finally {
+      fs.rmSync(dir, { recursive: true, force: true });
+    }
+  });
+
+  it('reports no-assembly before the app is synthesized', async () => {
+    const client = createTestClient({ readAssembly: async () => ({ status: 'not-found' }) });
+    await initializeClient(client, { applicationDir: '/p' });
+
+    const result = client.handlers.onGetConstructTree();
+
+    expect(result.status).toBe('no-assembly');
+    expect(result.entries).toEqual([]);
+    expect(result.assemblyDir).toBe(path.join('/p', 'cdk.out'));
+  });
+});