Fix global variable context resolution

package.json

@@ -110,6 +110,7 @@
             "original",
             "agent",
             "cfg",
+            "cfg_fallback",
             "experimental",
             "fastest",
             "best"

src/agents/contextSelector.ts

@@ -7,7 +7,7 @@ import { getDecodedTokensFromSymbol, countUniqueDefinitions, getTokensFromStr }
 import { DecodedToken } from '../lsp/types';
 import { retrieveDef, isBetweenFocalMethod, isInWorkspace } from '../lsp/definition';
 import { getReferenceInfo } from '../lsp/reference';
-import { formatToJSON, extractArrayFromJSON } from '../lsp/utils';
+import { formatToJSON, extractArrayFromJSON, extractHoverText, getHover } from '../lsp/utils';
 import { getSymbolDetail } from '../lsp/symbol';
 import { getSymbolByLocation } from '../lsp/symbol';
 import { ContextSelectorConfig, findTemplateFile } from '../prompts/promptBuilder';
@@ -636,9 +636,9 @@ export class ContextSelector {
             return false;
         }
 
-        if (token.type === 'variable' || token.type === 'property') {
+        if (this.isVariableLikeToken(token)) {
             console.log(`[addDefinitionToTerm] Processing ${token.type} definition for: ${term.name}`);
-            return this.addVariableDefinition(term, token, defSymbolDoc);
+            return await this.addVariableDefinition(term, token, defSymbolDoc);
         } else {
             console.log(`[addDefinitionToTerm] Processing method definition for: ${term.name}`);
             return this.addMethodDefinition(term, token, defSymbolDoc, functionSymbol);
@@ -669,11 +669,32 @@ export class ContextSelector {
             .join('\n');
     }
 
-    private addVariableDefinition(
+    private isVariableLikeToken(token: DecodedToken): boolean {
+        return ['variable', 'property', 'constant', 'member', 'global'].includes(token.type);
+    }
+
+    private async getHoverContextForToken(
+        token: DecodedToken,
+        defSymbolDoc: vscode.TextDocument
+    ): Promise<string | null> {
+        if (!token.definition[0].range) {
+            return null;
+        }
+
+        const defSymbol = token.defSymbol ?? await getSymbolByLocation(defSymbolDoc, token.definition[0].range.start);
+        if (!defSymbol) {
+            return null;
+        }
+
+        const hoverResults = await getHover(defSymbolDoc, defSymbol);
+        return extractHoverText(hoverResults);
+    }
+
+    private async addVariableDefinition(
         term: ContextTerm,
         token: DecodedToken,
         defSymbolDoc: vscode.TextDocument
-    ): boolean {
+    ): Promise<boolean> {
         console.log(`[addVariableDefinition] Starting for: ${term.name}`);
         const targetLine = defSymbolDoc.lineAt(token.definition[0].range!.start.line).text.trim();
         if (this.document.getText(this.targetSymbol.range).includes(targetLine)) {
@@ -689,10 +710,22 @@ export class ContextSelector {
         for (let i = startLine; i <= endLine; i++) {
             contextLines.push(defSymbolDoc.lineAt(i).text.trim());
         }
-        term.context = this.addLineNumbers(contextLines.join('\n'), startLine + 1);
+        const variableContext = contextLines.join('\n').trim();
+        if (variableContext.length > 0) {
+            term.context = this.addLineNumbers(variableContext, startLine + 1);
+            console.log(`[addVariableDefinition] Successfully added definition for: ${term.name}`);
+            return true;
+        }
+
+        const hoverText = await this.getHoverContextForToken(token, defSymbolDoc);
+        if (hoverText) {
+            term.context = this.addLineNumbers(hoverText, token.definition[0].range!.start.line + 1);
+            console.log(`[addVariableDefinition] Successfully added hover definition for: ${term.name}`);
+            return true;
+        }
 
-        console.log(`[addVariableDefinition] Successfully added definition for: ${term.name}`);
-        return true;
+        console.log(`[addVariableDefinition] Failed to add definition for: ${term.name}`);
+        return false;
     }
 
     private async addMethodDefinition(
@@ -732,4 +765,4 @@ export class ContextSelector {
 // Export a convenience function to get the singleton instance
 export async function getContextSelectorInstance(document: vscode.TextDocument, targetSymbol: vscode.DocumentSymbol): Promise<ContextSelector> {
     return await ContextSelector.getInstance(document, targetSymbol);
-}
+}

src/config.ts

@@ -36,6 +36,7 @@ export enum GenerationType {
     ORIGINAL = 'original', // with context, naive context, only template
     LSPRAG = 'lsprag',
     AGENT = 'agent',
+    CFG_FALLBACK = 'cfg_fallback',
     EXPERIMENTAL = 'experimental',
     FASTEST = 'fastest',
     BEST = 'best',

src/lsp/definition.ts

@@ -191,7 +191,6 @@ export function isBetweenFocalMethod(
     );
 }
 
-
 /** Normalize a URI string or filesystem path into an absolute fsPath. */
 function toFsPathNormalized(input: string): string | null {
   if (!input) return null;
@@ -223,11 +222,20 @@ export function isInWorkspace(uriOrPath: string): boolean {
     .map(f => toFsPathNormalized(f.uri.fsPath))
     .filter(Boolean) as string[];
 
-  // Prefix match under workspace root
-  return roots.some(root => target.startsWith(ensureTrailingSep(root)));
-}
+  const configuredWorkspace = toFsPathNormalized(getConfigInstance().workspace);
+  if (configuredWorkspace && !roots.includes(configuredWorkspace)) {
+    roots.push(configuredWorkspace);
+  }
 
+  if (roots.length === 0) {
+    return false;
+  }
 
+  return roots.some(root => {
+    const normalizedRoot = ensureTrailingSep(root);
+    return target === root || target.startsWith(normalizedRoot);
+  });
+}
 export async function classifyTokenByUri(document: vscode.TextDocument, DefUseMap: DecodedToken[], parentSymbol: vscode.DocumentSymbol | null = null): Promise<Map<string, DecodedToken[]>> {
     // Get all definitions from DefUseMap, but we retreive the method and definition together
     // Define the structure for ParentDefinition

src/lsp/utils.ts

@@ -141,6 +141,28 @@ export async function getHover(document: vscode.TextDocument, symbol: vscode.Doc
 }
 }
 
+export function extractHoverText(hover: vscode.Hover | undefined): string | null {
+    if (!hover) {
+        return null;
+    }
+
+    const parts = hover.contents.map(content => {
+        if (typeof content === 'string') {
+            return content;
+        }
+        if (content instanceof vscode.MarkdownString) {
+            return content.value;
+        }
+        return content.value;
+    }).filter(text => text && text.trim().length > 0);
+
+    if (parts.length === 0) {
+        return null;
+    }
+
+    return parts.join('\n').trim();
+}
+
 export function removeComments(code: string): string {
     const commentRegex = [
         /\/\/[^\n]*\n/g,
@@ -303,4 +325,4 @@ export function extractArrayFromJSON(jsonContent: any): any[] {
 function isPublic(symbol: vscode.DocumentSymbol, document: vscode.TextDocument): boolean {
     const funcDefinition = document.lineAt(symbol.selectionRange.start.line).text;
     return funcDefinition.includes('public') || false;
-}
+}

src/strategy/generators/cfgFallback.ts

@@ -0,0 +1,39 @@
+import { ConditionAnalysis } from '../../cfg/path';
+import { getContextTermsFromAllTokens } from '../../tokenAnalyzer';
+import { ContextSelector, ContextTerm } from '../../agents/contextSelector';
+import { getConfigInstance } from '../../config';
+import { getReferenceInfo } from '../../lsp/reference';
+import { saveContextTerms } from '../../fileHandler';
+import { LSPRAGTestGenerator } from './lsprag';
+
+export class CFGFallbackTestGenerator extends LSPRAGTestGenerator {
+    protected async collectInfo(
+        _conditions: ConditionAnalysis[] = [],
+        _functionInfo: Map<string, string> = new Map()
+    ): Promise<ContextTerm[] | null> {
+        let enrichedTerms: ContextTerm[] = [];
+        const tokenCollectTime = Date.now();
+        const contextSelector = await ContextSelector.create(this.document, this.functionSymbol);
+        const tokens = await contextSelector.loadTokens();
+        const identifiedTerms = await getContextTermsFromAllTokens(this.functionSymbol, tokens);
+        this.logger.log("getContextTermsFromAllTokens", (Date.now() - tokenCollectTime).toString(), null, "");
+
+        if (!await this.reportProgress(`[${getConfigInstance().generationType} mode] - gathering context`, 20)) {
+            return null;
+        }
+
+        const retrieveTime = Date.now();
+        enrichedTerms = await contextSelector.gatherContext(identifiedTerms, this.functionSymbol);
+        const referenceStrings = await getReferenceInfo(this.document, this.functionSymbol.selectionRange, 60, false);
+        const contextTermsForFunctionSymbol: ContextTerm = {
+            name: this.functionSymbol.name,
+            context: referenceStrings,
+            need_example: true,
+            hint: ["focal method"]
+        };
+        enrichedTerms.unshift(contextTermsForFunctionSymbol);
+        this.logger.log("gatherContext", (Date.now() - retrieveTime).toString(), null, "");
+        saveContextTerms(this.sourceCode, enrichedTerms, getConfigInstance().logSavePath!, this.fileName);
+        return enrichedTerms;
+    }
+}

src/strategy/generators/factory.ts

@@ -7,6 +7,7 @@ import { SymPromptTestGenerator } from './symPrompt';
 import { NaiveTestGenerator } from './naive';
 import { CFGTestGenerator } from './cfg';
 import { LSPRAGTestGenerator } from './lsprag';
+import { CFGFallbackTestGenerator } from './cfgFallback';
 
 // Factory to create the appropriate generator
 export function createTestGenerator(
@@ -30,6 +31,8 @@ export function createTestGenerator(
 			return new AgentTestGenerator(document, functionSymbol, languageId, fileName, logger, progress, token, srcPath);
 		case GenerationType.CFG:
 			return new CFGTestGenerator(document, functionSymbol, languageId, fileName, logger, progress, token, srcPath);
+		case GenerationType.CFG_FALLBACK:
+			return new CFGFallbackTestGenerator(document, functionSymbol, languageId, fileName, logger, progress, token, srcPath);
 		case GenerationType.LSPRAG:
 			return new LSPRAGTestGenerator(document, functionSymbol, languageId, fileName, logger, progress, token, srcPath);
 		default:

src/test/fixtures/python/global_constant.py

@@ -0,0 +1,17 @@
+c = 2
+
+
+class Calculator:
+    def add(self, x, y):
+        result = x + y
+        return result
+
+    def multiply(self, x, y):
+        result = x * y
+        return result
+
+    def complex_calculation(self, a, b):
+        intermediate_result = self.add(a, b)
+        intermediate_result = self.add(intermediate_result, c)
+
+        return intermediate_result

src/test/suite/lsp/globalConstant.test.ts

@@ -0,0 +1,44 @@
+import * as assert from 'assert';
+import * as vscode from 'vscode';
+import * as path from 'path';
+import { ContextSelector } from '../../../agents/contextSelector';
+import { GenerationType, getConfigInstance, PromptType } from '../../../config';
+import { setWorkspaceFolders } from '../../../helper';
+import { getSymbolFromDocument } from '../../../lsp/symbol';
+import { getContextTermsFromAllTokens } from '../../../tokenAnalyzer';
+
+suite('LSP-Features: Global Constant Context Test', () => {
+    const fixturesDir = path.join(__dirname, '../../../../src/test/fixtures');
+    const pythonProjectPath = path.join(fixturesDir, 'python');
+
+    test('Python - Global variable context fallback test', async function() {
+        getConfigInstance().updateConfig({
+            workspace: pythonProjectPath,
+            generationType: GenerationType.CFG_FALLBACK,
+            promptType: PromptType.WITHCONTEXT
+        });
+
+        const workspaceFolders = setWorkspaceFolders(pythonProjectPath);
+        console.log(`Python workspace path: ${workspaceFolders[0].uri.fsPath}`);
+
+        const fileUri = vscode.Uri.file(path.join(pythonProjectPath, 'global_constant.py'));
+        const document = await vscode.workspace.openTextDocument(fileUri);
+        const symbol = await getSymbolFromDocument(document, 'complex_calculation');
+        assert.ok(symbol, 'Should find complex_calculation function');
+
+        const contextSelector = await ContextSelector.create(document, symbol!);
+        const tokens = await contextSelector.loadTokens();
+        const globalToken = tokens.find(t => t.word === 'c');
+        assert.ok(globalToken, 'Should find global variable token c');
+
+        const contextTerms = await getContextTermsFromAllTokens(symbol!, tokens);
+        const globalTerm = contextTerms.find(term => term.name === 'c');
+        assert.ok(globalTerm, 'Should collect global variable term c');
+        assert.ok(globalTerm!.need_definition, 'Should request definition context for c');
+
+        const enrichedTerms = await contextSelector.gatherContext(contextTerms, symbol!);
+        const enrichedGlobalTerm = enrichedTerms.find(term => term.name === 'c');
+        assert.ok(enrichedGlobalTerm, 'Should enrich context for global variable c');
+        assert.ok(enrichedGlobalTerm!.context?.includes('c = 2'), 'Should include global variable definition in context');
+    });
+});

src/test/suite/lsp/token.test.ts

@@ -2,7 +2,6 @@ import * as assert from 'assert';
 import * as vscode from 'vscode';
 import * as path from 'path';
 import { getAllSymbols, getSymbolFromDocument } from '../../../lsp/symbol';
-import { getContextSelectorInstance } from '../../../agents/contextSelector';
 import { getConfigInstance, Provider, GenerationType, PromptType } from '../../../config';
 import { activate, setPythonExtraPaths } from '../../../lsp/helper';
 import { setWorkspaceFolders } from '../../../helper';
@@ -106,4 +105,3 @@ suite('LSP-Features: Token Finding Test', () => {
     });
 
 });
-

src/tokenAnalyzer.ts

@@ -276,6 +276,28 @@ async function cfgGetContextTermsFromTokens(
 function removeFocalMethodFromContextTerms(contextTerms: ContextTerm[], focalMethodName: string): ContextTerm[] {
     return contextTerms.filter(term => term.name !== focalMethodName);
 }
+
+export async function getContextTermsFromAllTokens(
+    symbol: vscode.DocumentSymbol,
+    tokens: DecodedToken[]
+): Promise<ContextTerm[]> {
+    const needContextTerms = await Promise.all(tokens.map(async token => ({
+        name: token.word,
+        need_definition: !needSkip(token) && await cfgBasedIsDefinitionHelpful(token),
+        need_example: !needSkip(token) && cfgBasedIsReferenceHelpful(token),
+        need_full_definition: isFunctionArg(token),
+        context: "",
+        example: "",
+        token: token,
+    })));
+
+    const filteredTerms = needContextTerms.filter(term => term.need_definition === true || term.need_example === true);
+    let uniqueTokens = removeRedundantTokens(filteredTerms);
+    uniqueTokens = removeFocalMethodFromContextTerms(uniqueTokens, symbol.name);
+    console.log("needContextTerms (fallback) :", uniqueTokens.map(term => [term.name, term.need_definition, term.need_example]));
+    return uniqueTokens;
+}
+
 // 4. Main exported functions delegate to the selected algorithm
 export async function getContextTermsFromTokens(
     document: vscode.TextDocument, 
@@ -297,4 +319,4 @@ export async function getContextTermsFromTokens(
         default:
             return defaultGetContextTermsFromTokens(tokens).filter(term => term.need_definition === true || term.need_example === true);
     }
-};
+};