feat: gather context

packages/salesforcedx-vscode-apex/src/commands/apexActionController.ts

@@ -11,7 +11,7 @@ import { URL } from 'url';
 import * as vscode from 'vscode';
 import { parse, stringify } from 'yaml';
 import { nls } from '../messages';
-import { ApexClassOASEligibleResponse, SymbolEligibility } from '../openApiUtilities/schemas';
+import { ApexClassOASEligibleResponse, ApexClassOASGatherContextResponse } from '../openApiUtilities/schemas';
 import { getTelemetryService } from '../telemetry/telemetry';
 import { MetadataOrchestrator } from './metadataOrchestrator';
 
@@ -28,6 +28,7 @@ export class ApexActionController {
       ? 'SFDX: Create Apex Action from This Class'
       : 'SFDX: Create Apex Action from Selected Method';
     let metadata;
+    let context;
     let name;
     const telemetryService = await getTelemetryService();
     try {
@@ -44,10 +45,14 @@ export class ApexActionController {
           if (!metadata) {
             throw new Error(nls.localize('extraction_failed', type));
           }
+          context = await this.metadataOrchestrator.gatherContext(sourceUri);
+          if (!context) {
+            throw new Error(nls.localize('cannot_gather_context'));
+          }
 
           // Step 3: Generate OpenAPI Document
           progress.report({ message: nls.localize('generate_openapi_document') });
-          const openApiDocument = await this.generateOpenAPIDocument(metadata);
+          const openApiDocument = await this.generateOpenAPIDocument(metadata, context);
 
           // Step 4: Write OpenAPI Document to File
           name = isClass ? path.basename(metadata.resourceUri, '.cls') : metadata?.symbols?.[0]?.docSymbol?.name;
@@ -87,14 +92,12 @@ export class ApexActionController {
    * @param metadata - The metadata of the methods.
    * @returns The OpenAPI document as a string.
    */
-  private generateOpenAPIDocument = async (metadata: ApexClassOASEligibleResponse): Promise<string> => {
+  private generateOpenAPIDocument = async (
+    metadata: ApexClassOASEligibleResponse,
+    context: ApexClassOASGatherContextResponse
+  ): Promise<string> => {
     const documentText = fs.readFileSync(new URL(metadata.resourceUri.toString()), 'utf8');
-    const className = path.basename(metadata.resourceUri, '.cls');
-    const methodNames = (metadata.symbols || [])
-      .filter((symbol: SymbolEligibility) => symbol.isApexOasEligible)
-      .map((symbol: SymbolEligibility) => symbol.docSymbol?.name)
-      .filter((name: string | undefined) => name);
-    const openAPIdocument = await this.metadataOrchestrator.sendPromptToLLM(documentText, methodNames, className);
+    const openAPIdocument = await this.metadataOrchestrator.sendPromptToLLM(documentText, context);
 
     // Convert the OpenAPI document to YAML
     return this.cleanupYaml(openAPIdocument);

packages/salesforcedx-vscode-apex/src/commands/metadataOrchestrator.ts

@@ -13,6 +13,7 @@ import {
   ApexClassOASEligibleRequest,
   ApexClassOASEligibleResponse,
   ApexClassOASEligibleResponses,
+  ApexClassOASGatherContextResponse,
   ApexOASEligiblePayload,
   ApexOASResource
 } from '../openApiUtilities/schemas';
@@ -98,6 +99,31 @@ export class MetadataOrchestrator {
     return response;
   };
 
+  public gatherContext = async (
+    sourceUri: vscode.Uri | vscode.Uri[]
+  ): Promise<ApexClassOASGatherContextResponse | undefined> => {
+    const telemetryService = await getTelemetryService();
+    let response;
+    const languageClient = languageClientUtils.getClientInstance();
+    if (languageClient) {
+      try {
+        response = (await languageClient?.sendRequest(
+          'apexoas/gatherContext',
+          sourceUri?.toString() ?? vscode.window.activeTextEditor?.document.uri.toString()
+        )) as ApexClassOASGatherContextResponse;
+        telemetryService.sendEventData('gatherContextSucceeded', { context: JSON.stringify(response) });
+      } catch (error) {
+        telemetryService.sendException(
+          'gatherContextFailed',
+          `${error} failed to send request to language server for ${path.basename(sourceUri.toString())}`
+        );
+        // fallback TBD after we understand it better
+        throw new Error(nls.localize('cannot_gather_context'));
+      }
+    }
+    return response;
+  };
+
   public validateEligibility = async (
     sourceUri: vscode.Uri | vscode.Uri[],
     isMethodSelected: boolean = false
@@ -155,7 +181,7 @@ export class MetadataOrchestrator {
       } else return ApexOASResource.class;
     }
   }
-  sendPromptToLLM = async (editorText: string, methods: string[], className: string): Promise<string> => {
+  sendPromptToLLM = async (editorText: string, context: ApexClassOASGatherContextResponse): Promise<string> => {
     console.log('This is the sendPromptToLLM() method');
     console.log('document text = ' + editorText);
 
@@ -174,7 +200,7 @@ export class MetadataOrchestrator {
       userPrompt +
       '\n\n***Code Context***\n```\n' +
       editorText +
-      `\nClass name: ${className}, methods: ${methods.join(',')}\n` +
+      `\nClass name: ${context.classDetail.name}, methods: ${context.methods.map(method => method.name).join(',')}\n` +
       `\n\`\`\`\n${endOfPromptTag}\n${assistantTag}`;
     console.log('input = ' + input);
     let result;

packages/salesforcedx-vscode-apex/src/messages/i18n.ja.ts

@@ -63,6 +63,7 @@ export const messages = {
     'Method %s is not eligible for Apex Action creation. It is not annotated with @AuraEnabled or has wrong access modifiers.',
   unknown: 'Unknown',
   invalid_active_text_editor: 'The active text editor is missing or is an invalid file.',
+  cannot_gather_context: 'An error occurred while gathering context for the Apex class.',
   sobjects_no_refresh_if_already_active_error_text:
     'sObject 定義の更新が既に実行中です。プロセスを再起動する必要がある場合は、実行中のタスクをキャンセルしてください。',
   test_view_loading_message: 'Apex テストを読み込んでいます...',

packages/salesforcedx-vscode-apex/src/messages/i18n.ts

@@ -106,5 +106,6 @@ export const messages = {
   orphan_process_advice:
     "The list of processes below are Apex Language Server instances that didn't properly shutdown. These\nprocesses can be stopped from the warning message that brought you here, or you can handle this\ntask yourself. If you choose to terminate these processes yourself, refer to relevant documentation\nto stop these processes.",
   unknown: 'Unknown',
-  invalid_active_text_editor: 'The active text editor is missing or is an invalid file.'
+  invalid_active_text_editor: 'The active text editor is missing or is an invalid file.',
+  cannot_gather_context: 'An error occurred while gathering context for the Apex class.'
 };

packages/salesforcedx-vscode-apex/src/openApiUtilities/schemas.ts

@@ -32,6 +32,46 @@ export type ApexOASEligiblePayload = {
 };
 export type ApexClassOASEligibleResponses = ApexClassOASEligibleResponse[];
 
+export type ApexClassOASGatherContextResponse = {
+  classDetail: ApexOASClassDetail;
+  properties: ApexOASPropertyDetail[];
+  methods: ApexOASMethodDetail[];
+  relationships: Map<string, Map<string, string[]>>; // Map<methodName, Map<srcClassUri, List<methodOrPropName>>>
+  documentations: Map<string, string[]>; // Map<method/prop/class name, each line of documentation>
+};
+
+export type ApexOASClassDetail = {
+  name: string;
+  interfaces: ApexOASInterface[];
+  extendedClass: ApexOASClassDetail | null;
+  annotations: string[];
+  definitionModifiers: string[];
+  accessModifiers: string[];
+  innerClasses: DocumentSymbol[];
+};
+
+export type ApexOASPropertyDetail = {
+  name: string;
+  type: string;
+  documentSymbol: DocumentSymbol;
+  modifiers: string[];
+  annotations: string[];
+};
+
+export type ApexOASMethodDetail = {
+  name: string;
+  returnType: string;
+  parameterTypes: string[];
+  modifiers: string[];
+  annotations: string[];
+};
+
+export type ApexOASInterface = {
+  name: string;
+  uri: string;
+  methods: DocumentSymbol[];
+};
+
 export enum ApexOASResource {
   class = 'CLASS',
   multiClass = 'MULTI CLASSES',

packages/salesforcedx-vscode-apex/test/jest/commands/metadataOrchestrator.test.ts

@@ -6,8 +6,10 @@
  */
 import { notificationService } from '@salesforce/salesforcedx-utils-vscode';
 import * as vscode from 'vscode';
+import { ApexLanguageClient } from '../../../src/apexLanguageClient';
 import { MetadataOrchestrator } from '../../../src/commands/metadataOrchestrator';
 import { languageClientUtils } from '../../../src/languageUtils';
+import { nls } from '../../../src/messages';
 import { ApexOASResource } from '../../../src/openApiUtilities/schemas';
 import { getTelemetryService } from '../../../src/telemetry/telemetry';
 import { MockTelemetryService } from '../telemetry/mockTelemetryService';
@@ -82,6 +84,52 @@ describe('MetadataOrchestrator', () => {
     });
   });
 
+  describe('gatherContext', () => {
+    let getClientInstanceSpy;
+    let mockTelemetryService: any;
+
+    beforeEach(() => {
+      (getTelemetryService as jest.Mock).mockResolvedValue(new MockTelemetryService());
+    });
+
+    afterEach(() => {
+      jest.clearAllMocks();
+    });
+
+    it('should send a request and return the response when successful', async () => {
+      const mockLanguageClient = {
+        sendRequest: jest.fn().mockResolvedValue({ some: 'response' })
+      } as unknown as ApexLanguageClient;
+
+      getClientInstanceSpy = jest.spyOn(languageClientUtils, 'getClientInstance').mockReturnValue(mockLanguageClient);
+
+      const mockUri = { path: '/hello/world.cls' } as vscode.Uri;
+      const response = await orchestrator.gatherContext(mockUri);
+
+      expect(mockLanguageClient.sendRequest).toHaveBeenCalledWith('apexoas/gatherContext', mockUri.toString());
+      expect(response).toEqual({ some: 'response' });
+    });
+
+    it('should handle language client being unavailable', async () => {
+      jest.spyOn(languageClientUtils, 'getClientInstance').mockReturnValue(undefined);
+
+      const response = await orchestrator.gatherContext(vscode.Uri.file('/path/to/source'));
+      expect(response).toBeUndefined();
+    });
+
+    it('should handle errors and throw a localized error', async () => {
+      const mockLanguageClient = {
+        sendRequest: jest.fn().mockRejectedValue(new Error('Some error'))
+      } as unknown as ApexLanguageClient;
+
+      jest.spyOn(languageClientUtils, 'getClientInstance').mockReturnValue(mockLanguageClient);
+
+      const mockUri = { path: '/hello/world.cls' } as vscode.Uri;
+
+      expect(orchestrator.gatherContext(mockUri)).rejects.toThrow(nls.localize('cannot_gather_context'));
+    });
+  });
+
   describe('validateEligibility', () => {
     let eligibilityDelegateSpy: jest.SpyInstance;
 
@@ -161,7 +209,6 @@ describe('MetadataOrchestrator', () => {
   });
 
   describe('eligibilityDelegate', () => {
-    let mockLanguageClient;
     let getClientInstanceSpy;
     beforeEach(() => {
       (getTelemetryService as jest.Mock).mockResolvedValue(new MockTelemetryService());