feat: gather context (#5988)

* feat: gather context * chore: update jar * chore: eligibility criteria for class and method annotations (#5981) * chore: eligibility criteria for class and method annotations * chore: update rest-related annotations * feat: latest jorje jar * chore: pin vsp version * feat: gather context * chore: update jar * chore: reorder requests * chore: reorder imports * chore: renaming after pr comments --------- Co-authored-by: Mingxuan Zhang <[email protected]>
forcedotcom · Jan 2, 2025 · 89b0928 · 89b0928
1 parent 420f306
commit 89b0928
Show file tree

Hide file tree

Showing 7 changed files with 155 additions and 33 deletions.
diff --git a/packages/salesforcedx-vscode-apex/out/apex-jorje-lsp.jar b/packages/salesforcedx-vscode-apex/out/apex-jorje-lsp.jar
diff --git a/packages/salesforcedx-vscode-apex/src/commands/apexActionController.ts b/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';
 
@@ -27,7 +27,8 @@ export class ApexActionController {
     const command = isClass
       ? 'SFDX: Create Apex Action from This Class'
       : 'SFDX: Create Apex Action from Selected Method';
-    let metadata;
+    let eligibilityResult;
+    let context;
     let name;
     const telemetryService = await getTelemetryService();
     try {
@@ -38,19 +39,27 @@ export class ApexActionController {
           cancellable: true
         },
         async progress => {
-          // Step 1: Extract Metadata
-          progress.report({ message: nls.localize('extract_metadata') });
-          metadata = await this.metadataOrchestrator.extractMetadata(sourceUri, !isClass);
-          if (!metadata) {
-            throw new Error(nls.localize('extraction_failed', type));
+          // Step 1: Validate eligibility
+          progress.report({ message: nls.localize('validate_eligibility') });
+          eligibilityResult = await this.metadataOrchestrator.validateMetadata(sourceUri, !isClass);
+          if (!eligibilityResult) {
+            throw new Error(nls.localize('class_validation_failed', type));
+          }
+
+          // Step 2: Gather context
+          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(eligibilityResult, context);
 
           // Step 4: Write OpenAPI Document to File
-          name = isClass ? path.basename(metadata.resourceUri, '.cls') : metadata?.symbols?.[0]?.docSymbol?.name;
+          name = isClass
+            ? path.basename(eligibilityResult.resourceUri, '.cls')
+            : eligibilityResult?.symbols?.[0]?.docSymbol?.name;
           const openApiFileName = `${name}_openapi.yml`;
           progress.report({ message: nls.localize('write_openapi_document_to_file') });
           await this.saveAndOpenDocument(openApiFileName, openApiDocument);
@@ -87,14 +96,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);

diff --git a/packages/salesforcedx-vscode-apex/src/commands/metadataOrchestrator.ts b/packages/salesforcedx-vscode-apex/src/commands/metadataOrchestrator.ts
@@ -13,6 +13,7 @@ import {
   ApexClassOASEligibleRequest,
   ApexClassOASEligibleResponse,
   ApexClassOASEligibleResponses,
+  ApexClassOASGatherContextResponse,
   ApexOASEligiblePayload,
   ApexOASResource
 } from '../openApiUtilities/schemas';
@@ -45,10 +46,10 @@ export interface Parameter {
  */
 export class MetadataOrchestrator {
   /**
-   * Extracts metadata for the method at the current cursor position.
+   * Validates and extracts metadata for the method at the current cursor position.
    * @returns The metadata of the method, or undefined if no method is found.
    */
-  public extractMetadata = async (
+  public validateMetadata = async (
     sourceUri: vscode.Uri | vscode.Uri[],
     isMethodSelected: boolean = false
   ): Promise<ApexClassOASEligibleResponse | undefined> => {
@@ -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;

diff --git a/packages/salesforcedx-vscode-apex/src/messages/i18n.ja.ts b/packages/salesforcedx-vscode-apex/src/messages/i18n.ja.ts
@@ -52,9 +52,9 @@ export const messages = {
   apex_test_run_description_text: 'Apex テストを実行',
   apex_test_run_text: 'SFDX: Apex テストを呼び出す',
   create_apex_action_failed: 'Failed to create Apex Action',
-  extract_metadata: 'Extracting metadata.',
-  extraction_failed: 'Failed to extract metadata from %s',
-  validation_failed: 'Failed to validate metadata.',
+  validate_eligibility: 'Validating eligibility.',
+  class_validation_failed: 'Failed to validate eligibility from %s',
+  validation_failed: 'Failed to validate eligibility.',
   apex_class_not_valid: 'The Apex Class %s is not valid for Open AI document generation.',
   apex_action_created: 'Apex Action created for %s: %s.',
   generate_openapi_document: 'Generating OpenAPI document.',
@@ -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 テストを読み込んでいます...',

diff --git a/packages/salesforcedx-vscode-apex/src/messages/i18n.ts b/packages/salesforcedx-vscode-apex/src/messages/i18n.ts
@@ -55,9 +55,9 @@ export const messages = {
   apex_test_run_codeAction_no_method_test_param_text:
     'Test method not provided. Run the code action on a method annotated with @isTest or testMethod.',
   create_apex_action_failed: 'Failed to create Apex Action',
-  extract_metadata: 'Extracting metadata.',
-  extraction_failed: 'Failed to extract metadata from %s',
-  validation_failed: 'Failed to validate metadata.',
+  validate_eligibility: 'Validating eligibility.',
+  class_validation_failed: 'Failed to validate eligibility from %s',
+  validation_failed: 'Failed to validate eligibility.',
   apex_class_not_valid: 'The Apex Class %s is not valid for Open AI document generation.',
   apex_action_created: 'Apex Action created for %s: %s.',
   generate_openapi_document: 'Generating OpenAPI document.',
@@ -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.'
 };
diff --git a/packages/salesforcedx-vscode-apex/src/openApiUtilities/schemas.ts b/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',

diff --git a/packages/salesforcedx-vscode-apex/test/jest/commands/metadataOrchestrator.test.ts b/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';
@@ -45,8 +47,8 @@ describe('MetadataOrchestrator', () => {
     });
     it('should throw an error if no eligible responses are returned', async () => {
       jest.spyOn(orchestrator, 'validateEligibility').mockResolvedValue(undefined);
-      await expect(orchestrator.extractMetadata(editorStub.document.uri)).rejects.toThrow(
-        'Failed to validate metadata.'
+      await expect(orchestrator.validateMetadata(editorStub.document.uri)).rejects.toThrow(
+        'Failed to validate eligibility.'
       );
     });
 
@@ -55,15 +57,15 @@ describe('MetadataOrchestrator', () => {
         { isApexOasEligible: false, isEligible: false, symbols: [{ docSymbol: { name: 'someMethod' } }] }
       ];
       jest.spyOn(orchestrator, 'validateEligibility').mockResolvedValue(mockResponse);
-      await expect(orchestrator.extractMetadata(editorStub.document.uri, true)).rejects.toThrow(
+      await expect(orchestrator.validateMetadata(editorStub.document.uri, true)).rejects.toThrow(
         'Method someMethod is not eligible for Apex Action creation. It is not annotated with @AuraEnabled or has wrong access modifiers.'
       );
     });
 
     it('should throw an error if the first eligible response is not eligible and method is not selected', async () => {
       const mockResponse: any = [{ isApexOasEligible: false, isEligible: false, resourceUri: '/hello/world.cls' }];
       jest.spyOn(orchestrator, 'validateEligibility').mockResolvedValue(mockResponse);
-      await expect(orchestrator.extractMetadata(editorStub.document.uri)).rejects.toThrow(
+      await expect(orchestrator.validateMetadata(editorStub.document.uri)).rejects.toThrow(
         'The Apex Class world is not valid for Open AI document generation.'
       );
     });
@@ -77,11 +79,57 @@ describe('MetadataOrchestrator', () => {
         }
       ];
       jest.spyOn(orchestrator, 'validateEligibility').mockResolvedValue(mockResponse);
-      const result = await orchestrator.extractMetadata(editorStub.document.uri);
+      const result = await orchestrator.validateMetadata(editorStub.document.uri);
       expect(result).toEqual(mockResponse[0]);
     });
   });
 
+  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());