feat: Enhance the context returned by the context provider (#1518)

* feat: support new command from deps

* chore: update

* chore: update

* feat: support deps command to enhance the context provider

* fix: update

Author: wenytang-ms

src/copilot/context/copilotHelper.ts

@@ -9,6 +9,10 @@ export interface INodeImportClass {
     uri: string;
     className: string;  // Changed from 'class' to 'className' to match Java code
 }
+
+export interface IProjectDependency {
+    [key: string]: string;
+}
 /**
  * Helper class for Copilot integration to analyze Java project dependencies
  */
@@ -23,8 +27,9 @@ export namespace CopilotHelper {
         if (cancellationToken?.isCancellationRequested) {
             return [];
         }
+
         // Ensure the Java Dependency extension is installed and meets the minimum version requirement.
-        if (!await validateExtensionInstalled("vscjava.vscode-java-dependency", "0.26.0")) {
+        if (!await validateExtensionInstalled("vscjava.vscode-java-dependency", "0.26.2")) {
             return [];
         }
 
@@ -43,11 +48,23 @@ export namespace CopilotHelper {
                         cancellationToken.onCancellationRequested(() => {
                             reject(new Error('Operation cancelled'));
                         });
+                    }),
+                    new Promise<INodeImportClass[]>((_, reject) => {
+                        setTimeout(() => {
+                            reject(new Error('Operation timed out'));
+                        }, 80); // 80ms timeout
                     })
                 ]);
                 return result || [];
             } else {
-                const result = await commandPromise;
+                const result = await Promise.race([
+                    commandPromise,
+                    new Promise<INodeImportClass[]>((_, reject) => {
+                        setTimeout(() => {
+                            reject(new Error('Operation timed out'));
+                        }, 80); // 80ms timeout
+                    })
+                ]);
                 return result || [];
             }
         } catch (error: any) {
@@ -59,4 +76,64 @@ export namespace CopilotHelper {
             return [];
         }
     }
+
+    /**
+     * Resolves project dependencies for the given project URI
+     * @param projectUri The URI of the Java project to analyze
+     * @param cancellationToken Optional cancellation token to abort the operation
+     * @returns Object containing project dependencies as key-value pairs
+     */
+    export async function resolveProjectDependencies(projectUri: Uri, cancellationToken?: CancellationToken): Promise<IProjectDependency> {
+        if (cancellationToken?.isCancellationRequested) {
+            return {};
+        }
+
+        // Ensure the Java Dependency extension is installed and meets the minimum version requirement.
+        if (!await validateExtensionInstalled("vscjava.vscode-java-dependency", "0.26.2")) {
+            return {};
+        }
+        
+        if (cancellationToken?.isCancellationRequested) {
+            return {};
+        }
+        
+        try {
+            // Create a promise that can be cancelled
+            const commandPromise = commands.executeCommand("java.execute.workspaceCommand", "java.project.getDependencies", projectUri.toString()) as Promise<IProjectDependency>;
+            //set timeout
+            if (cancellationToken) {
+                const result = await Promise.race([
+                    commandPromise,
+                    new Promise<IProjectDependency>((_, reject) => {
+                        cancellationToken.onCancellationRequested(() => {
+                            reject(new Error('Operation cancelled'));
+                        });
+                    }),
+                    new Promise<IProjectDependency>((_, reject) => {
+                        setTimeout(() => {
+                            reject(new Error('Operation timed out'));
+                        }, 40); // 40ms timeout
+                    })
+                ]);
+                return result || {};
+            } else {
+                const result = await Promise.race([
+                    commandPromise,
+                    new Promise<IProjectDependency>((_, reject) => {
+                        setTimeout(() => {
+                            reject(new Error('Operation timed out'));
+                        }, 40); // 40ms timeout
+                    })
+                ]);
+                return result || {};
+            }
+        } catch (error: any) {
+            if (error.message === 'Operation cancelled') {
+                logger.info('Resolve project dependencies cancelled');
+                return {};
+            }
+            logger.error("Error resolving project dependencies:", error);
+            return {};
+        }
+    }
 }

src/copilot/contextProvider.ts

@@ -11,7 +11,7 @@ import * as vscode from 'vscode';
 import { CopilotHelper } from './context/copilotHelper';
 import { sendInfo } from "vscode-extension-telemetry-wrapper";
 import {
-    logger, 
+    logger,
     JavaContextProviderUtils,
     CancellationError,
     InternalCancellationError,
@@ -93,13 +93,15 @@ function createJavaContextResolver(): ContextResolverFunction {
 /**
  * Send telemetry data for Java context resolution
  */
-function sendContextTelemetry(request: ResolveRequest, start: number, itemCount: number, status: string, error?: string) {
+function sendContextTelemetry(request: ResolveRequest, start: number, items: SupportedContextItem[], status: string, error?: string) {
     const duration = Math.round(performance.now() - start);
+    const tokenCount = JavaContextProviderUtils.calculateTokenCount(items);
     const telemetryData: any = {
         "action": "resolveJavaContext",
         "completionId": request.completionId,
         "duration": duration,
-        "itemCount": itemCount,
+        "itemCount": items.length,
+        "tokenCount": tokenCount,
         "status": status
     };
 
@@ -128,16 +130,37 @@ async function resolveJavaContext(request: ResolveRequest, copilotCancel: vscode
 
         const document = activeEditor.document;
 
+        // Resolve project dependencies first
+        const projectDependencies = await CopilotHelper.resolveProjectDependencies(document.uri, copilotCancel);
+        logger.info('Resolved project dependencies count:', Object.keys(projectDependencies).length);
+        
+        // Check for cancellation after dependency resolution
+        JavaContextProviderUtils.checkCancellation(copilotCancel);
+        
+        // Convert project dependencies to Trait items
+        if (projectDependencies && Object.keys(projectDependencies).length > 0) {
+            for (const [key, value] of Object.entries(projectDependencies)) {
+                items.push({
+                    name: key,
+                    value: value,
+                    importance: 50
+                });
+            }
+        }
+        
+        // Check for cancellation before resolving imports
+        JavaContextProviderUtils.checkCancellation(copilotCancel);
+
         // Resolve imports directly without caching
         const importClass = await CopilotHelper.resolveLocalImports(document.uri, copilotCancel);
-        logger.trace('Resolved imports count:', importClass?.length || 0);
+        logger.info('Resolved imports count:', importClass?.length || 0);
 
         // Check for cancellation after resolution
         JavaContextProviderUtils.checkCancellation(copilotCancel);
 
         // Check for cancellation before processing results
         JavaContextProviderUtils.checkCancellation(copilotCancel);
-        
+
         if (importClass) {
             // Process imports in batches to reduce cancellation check overhead
             const contextItems = JavaContextProviderUtils.createContextItemsFromImports(importClass);
@@ -149,26 +172,26 @@ async function resolveJavaContext(request: ResolveRequest, copilotCancel: vscode
         }
     } catch (error: any) {
         if (error instanceof CopilotCancellationError) {
-            sendContextTelemetry(request, start, items.length, "cancelled_by_copilot");
+            sendContextTelemetry(request, start, items, "cancelled_by_copilot");
             throw error;
         }
         if (error instanceof vscode.CancellationError || error.message === CancellationError.Canceled) {
-            sendContextTelemetry(request, start, items.length, "cancelled_internally");
+            sendContextTelemetry(request, start, items, "cancelled_internally");
             throw new InternalCancellationError();
         }
 
         // Send telemetry for general errors (but continue with partial results)
-        sendContextTelemetry(request, start, items.length, "error_partial_results", error.message || "unknown_error");
+        sendContextTelemetry(request, start, items, "error_partial_results", error.message || "unknown_error");
 
         logger.error(`Error resolving Java context for ${documentUri}:${caretOffset}:`, error);
 
         // Return partial results and log completion for error case
         JavaContextProviderUtils.logCompletion('Java context resolution', documentUri, caretOffset, start, items.length);
         return items;
     }
-    
+
     // Send telemetry data once at the end for success case
-    sendContextTelemetry(request, start, items.length, "succeeded");
+    sendContextTelemetry(request, start, items, "succeeded");
 
     JavaContextProviderUtils.logCompletion('Java context resolution', documentUri, caretOffset, start, items.length);
     return items;

src/copilot/utils.ts

@@ -301,6 +301,39 @@ export class JavaContextProviderUtils {
 
         return installCount;
     }
+
+    /**
+     * Calculate approximate token count for context items
+     * Using a simple heuristic: ~4 characters per token
+     * Optimized for performance by using reduce and direct property access
+     */
+    static calculateTokenCount(items: SupportedContextItem[]): number {
+        // Fast path: if no items, return 0
+        if (items.length === 0) {
+            return 0;
+        }
+        
+        // Use reduce for better performance
+        const totalChars = items.reduce((sum, item) => {
+            let itemChars = 0;
+            // Direct property access is faster than 'in' operator
+            const value = (item as any).value;
+            const name = (item as any).name;
+            
+            if (value && typeof value === 'string') {
+                itemChars += value.length;
+            }
+            if (name && typeof name === 'string') {
+                itemChars += name.length;
+            }
+            
+            return sum + itemChars;
+        }, 0);
+        
+        // Approximate: 1 token ≈ 4 characters
+        // Use bitwise shift for faster division by 4
+        return (totalChars >> 2) + (totalChars & 3 ? 1 : 0);
+    }
 }
 
 /**