diff --git a/__mocks__/vscode.ts b/__mocks__/vscode.ts index 32f693665..a06c62a0c 100644 --- a/__mocks__/vscode.ts +++ b/__mocks__/vscode.ts @@ -129,18 +129,28 @@ export class DiagnosticCollection } export class EventEmitter { - listeners = new Set<(...args: any[]) => void>(); - event = (listener: (e: T) => any) => { - this.listeners.add(listener); - return () => { - this.listeners.delete(listener); - }; - }; + private listeners = new Set<(data: T) => any>(); + + event = vi + .fn() + .mockName('event') + .mockImplementation((listener: (e: T) => any) => { + this.listeners.add(listener); + return { + dispose: vi + .fn() + .mockName('dispose') + .mockImplementation(() => { + this.listeners.delete(listener); + }), + }; + }); + fire = vi .fn() .mockName('fire') - .mockImplementation((event: T) => { - this.listeners.forEach(listener => listener(event)); + .mockImplementation((data: T): void => { + this.listeners.forEach(listener => listener(data)); }); } @@ -354,7 +364,10 @@ export class Uri { static joinPath = vi .fn() .mockName('joinPath') - .mockImplementation((...args) => Uri.parse(args.join('/'))); + .mockImplementation((...args) => { + const filteredArgs = args.filter(a => a.toString().length > 0); + return Uri.parse(filteredArgs.join('/')); + }); static parse = vi .fn() diff --git a/docs/python-remote-file-sourcing.md b/docs/python-remote-file-sourcing.md index 4be24f1a9..b31967244 100644 --- a/docs/python-remote-file-sourcing.md +++ b/docs/python-remote-file-sourcing.md @@ -156,3 +156,74 @@ def dashboard_content(table): 2. In your VS Code workspace, use the Deephaven extension to run `main.py`. The imports will resolve because the `stock_ticker` folder is registered as a remote file source. ![Run main.py](assets/run-main-py.gif) + +## Controller Import Prefix Support (Enterprise) + +**Deephaven Enterprise** uses a controller import registration mechanism (`meta_import()`) that requires modules to be importable under a prefixed name (e.g., `controller.mymodule`) for the server to find them. The VS Code extension automatically detects this pattern in your Python code and registers both the unprefixed and prefixed names with the server for each folder marked as a remote file source, allowing prefixed imports to be sourced by the extension. + +### Auto-Detection + +When you run Python code, the extension scans for `meta_import()` calls and infers the prefix to use. No extra setup is required — if your code already calls `meta_import()`, the extension picks it up automatically. + +**With default prefix (`controller`):** + +```python +import deephaven_enterprise.controller_import + +deephaven_enterprise.controller_import.meta_import() +``` + +**With a custom prefix:** + +```python +import deephaven_enterprise.controller_import + +deephaven_enterprise.controller_import.meta_import("myprefix") +``` + +**From-import style:** + +```python +from deephaven_enterprise.controller_import import meta_import + +meta_import("custom") +``` + +### Behavior + +- For each folder marked as a remote file source, both the unprefixed and prefixed module names are registered with the Deephaven server. +- Example: If you mark a folder called `mymodule` and a prefix of `controller` is detected, the server will receive both `mymodule` and `controller.mymodule` as importable names for that folder. +- Without a detected or configured prefix, only the unprefixed name (`mymodule`) is registered for each marked folder. +- Prefixes are **updated when you run Python code**: running a full file replaces any previous prefixes; running a snippet updates prefixes only if a `meta_import()` call is found in that snippet. + +### Manual Override + +You can set one or more prefixes explicitly in your VS Code settings: + +```json +"deephaven.importPrefixes": ["controller"] +``` + +When set, this array is used as the source of truth and auto-detection is skipped entirely. Multiple prefixes can be provided if needed: + +```json +"deephaven.importPrefixes": ["controller", "custom"] +``` + +The main reasons to set this manually are: + +- **Unrecognized imports** — the auto-detection doesn't recognize all possible import patterns such as aliasing: + + ```python + import deephaven_enterprise.controller_import as ci + + ci.meta_import() + + from deephaven_enterprise.controller_import import meta_import as m + + m() + ``` + +- **`meta_import()` in a dependency** — auto-detection only scans the code being run directly, not modules it imports. If the registration happens inside a dependency rather than the script itself, the prefix will not be detected. + +If either case applies, set `deephaven.importPrefixes` manually. diff --git a/package.json b/package.json index d4e5d9e44..1e7374e9b 100644 --- a/package.json +++ b/package.json @@ -160,6 +160,17 @@ }, "default": [] }, + "deephaven.importPrefixes": { + "type": [ + "array", + "null" + ], + "default": null, + "items": { + "type": "string" + }, + "markdownDescription": "Optional controller import prefixes for remote file sourcing. When set, overrides automatic detection from `meta_import()` calls in Python code. Defaults to automatic detection if not specified." + }, "deephaven.mcp.enabled": { "type": "boolean", "default": false, diff --git a/src/common/constants.ts b/src/common/constants.ts index cf55f7f46..3506aa0ce 100644 --- a/src/common/constants.ts +++ b/src/common/constants.ts @@ -16,6 +16,7 @@ export const CONFIG_KEY = { root: 'deephaven', coreServers: 'coreServers', enterpriseServers: 'enterpriseServers', + importPrefixes: 'importPrefixes', mcpAutoUpdateConfig: 'mcp.autoUpdateConfig', mcpDocsEnabled: 'mcp.docsEnabled', mcpEnabled: 'mcp.enabled', diff --git a/src/services/ConfigService.spec.ts b/src/services/ConfigService.spec.ts index 9a90e4e6d..d034e99b2 100644 --- a/src/services/ConfigService.spec.ts +++ b/src/services/ConfigService.spec.ts @@ -353,3 +353,39 @@ describe('updateWindsurfMcpConfig', () => { expect(getEnsuredContent).toHaveBeenCalledWith(windowsConfigUri, '{}\n'); }); }); + +describe('getImportPrefixes', () => { + it.each([ + ['not configured', undefined, undefined], + ['empty array', [], []], + ['single valid prefix', ['controller'], ['controller']], + ['multiple valid prefixes', ['controller', '_my_prefix'], ['controller', '_my_prefix']], + ])('should return without errors when %s', (_label, given, expected) => { + configMap.set(CONFIG_KEY.importPrefixes, given); + + const result = ConfigService.getImportPrefixes(); + + expect(result).toEqual(expected); + expect(vscode.window.showErrorMessage).not.toHaveBeenCalled(); + }); + + it.each([ + ['one invalid entry (empty string)', [''], [], ['']], + ['one invalid entry (starts with digit)', ['1controller'], [], ['1controller']], + ['one invalid entry (contains hyphen)', ['my-prefix'], [], ['my-prefix']], + ['one invalid entry (dotted name)', ['my.prefix'], [], ['my.prefix']], + ['mixed valid and invalid entries', ['valid', '1bad', 'also_valid', 'my-prefix'], ['valid', 'also_valid'], ['1bad', 'my-prefix']], + ])('should filter invalid entries and show errors when %s', (_label, given, expectedResult, expectedInvalid) => { + configMap.set(CONFIG_KEY.importPrefixes, given); + + const result = ConfigService.getImportPrefixes(); + + expect(result).toEqual(expectedResult); + expect(vscode.window.showErrorMessage).toHaveBeenCalledTimes(expectedInvalid.length); + for (const invalid of expectedInvalid) { + expect(vscode.window.showErrorMessage).toHaveBeenCalledWith( + expect.stringContaining(`'${invalid}' is not a valid import prefix name`) + ); + } + }); +}); diff --git a/src/services/ConfigService.ts b/src/services/ConfigService.ts index 0c0a1b675..bc9019f32 100644 --- a/src/services/ConfigService.ts +++ b/src/services/ConfigService.ts @@ -78,6 +78,28 @@ function hasValidURL({ url }: { url: string }): boolean { } } +// ASCII subset of Python identifier rules (PEP 3131 / py3 lexical spec allows Unicode, +// but controller prefix names are expected to be ASCII in practice). A prefix is a +// single identifier, not a dotted module path. +const VALID_PYTHON_IDENTIFIER_RE = /^[A-Za-z_][A-Za-z0-9_]*$/; + +function getImportPrefixes(): string[] | undefined { + const config = getConfig().get(CONFIG_KEY.importPrefixes); + if (config == null) { + return undefined; + } + + return config.filter(prefix => { + if (VALID_PYTHON_IDENTIFIER_RE.test(prefix)) { + return true; + } + vscode.window.showErrorMessage( + `Invalid 'deephaven.importPrefixes' setting: '${prefix}' is not a valid import prefix name. It will be ignored.` + ); + return false; + }); +} + async function toggleMcp(enable?: boolean): Promise { const currentState = isMcpEnabled(); const targetState = enable ?? !currentState; @@ -233,6 +255,7 @@ export async function updateWindsurfMcpConfig( export const ConfigService: IConfigService = { getCoreServers, getEnterpriseServers, + getImportPrefixes, isElectronFetchEnabled, isMcpDocsEnabled, isMcpEnabled, diff --git a/src/services/DhcService.spec.ts b/src/services/DhcService.spec.ts new file mode 100644 index 000000000..4adcbd017 --- /dev/null +++ b/src/services/DhcService.spec.ts @@ -0,0 +1,152 @@ +import * as vscode from 'vscode'; +import type { dh as DhcType } from '@deephaven/jsapi-types'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { DhcService } from './DhcService'; +import { ConfigService } from './ConfigService'; +import type { RemoteFileSourceService } from './RemoteFileSourceService'; +import type { UniqueID } from '../types'; + +vi.mock('vscode'); + +vi.mock('./ConfigService', () => { + const mockConfigService = { + getImportPrefixes: vi.fn(), + }; + // eslint-disable-next-line @typescript-eslint/naming-convention + return { ConfigService: mockConfigService }; +}); + +/** Build a minimal DhcService instance that has a session already set up. */ +function createTestDhcService({ + pythonRemoteFileSourcePlugin, + setControllerImportPrefixes = vi.fn(), + setPythonServerExecutionContext = vi.fn().mockResolvedValue(undefined), + sessionRunCode = vi.fn().mockResolvedValue({ + error: '', + changes: { created: [], updated: [], removed: [] }, + }), +}: { + pythonRemoteFileSourcePlugin?: DhcType.Widget | null; + setControllerImportPrefixes?: ReturnType; + setPythonServerExecutionContext?: ReturnType; + sessionRunCode?: ReturnType; +} = {}): DhcService { + const mockSession = { + runCode: sessionRunCode, + } as unknown as DhcType.IdeSession; + + const mockCn = { + getConsoleTypes: vi.fn().mockResolvedValue(['python']), + } as unknown as DhcType.IdeConnection; + + const mockRemoteFileSourceService = { + setControllerImportPrefixes, + setPythonServerExecutionContext, + } as unknown as RemoteFileSourceService; + + const mockDiagnosticsCollection = { + set: vi.fn(), + clear: vi.fn(), + } as unknown as vscode.DiagnosticCollection; + + const service = Object.assign(Object.create(DhcService.prototype), { + session: mockSession, + cn: mockCn, + cnId: 'test-cn-id' as UniqueID, + pythonRemoteFileSourcePlugin: + pythonRemoteFileSourcePlugin !== undefined + ? pythonRemoteFileSourcePlugin + : ({} as DhcType.Widget), + groovyRemoteFileSourcePluginService: null, + remoteFileSourceService: mockRemoteFileSourceService, + diagnosticsCollection: mockDiagnosticsCollection, + groovyDiagnosticsCollection: mockDiagnosticsCollection, + outputChannel: { + appendLine: vi.fn(), + show: vi.fn(), + } as unknown as vscode.OutputChannel, + toaster: { error: vi.fn(), info: vi.fn() }, + _isRunningCode: false, + _onDidChangeRunningCodeStatus: { fire: vi.fn() }, + disposables: { add: vi.fn() }, + serverUrl: new URL('http://localhost:10000/'), + }); + + return service; +} + +function mockTextDoc(codeText: string): vscode.TextDocument { + return { + uri: vscode.Uri.file('/path/to/file.py'), + getText: vi.fn().mockReturnValue(codeText), + } as unknown as vscode.TextDocument; +} + +describe('DhcService.runCode – importPrefixes setting', () => { + const mockSetControllerImportPrefixes = vi.fn(); + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it.each([ + { + label: 'uses setting prefixes when configured', + configPrefixes: ['myPrefix'], + input: 'x = 1', + expected: new Set(['myPrefix']), + }, + { + label: 'uses all prefixes when multiple configured', + configPrefixes: ['prefix1', 'prefix2'], + input: 'x = 1', + expected: new Set(['prefix1', 'prefix2']), + }, + { + label: 'falls back to extraction when undefined and code has prefixes', + configPrefixes: undefined, + input: 'deephaven_enterprise.controller_import.meta_import("myPrefix")\n', + expected: new Set(['myPrefix']), + }, + { + label: 'skips call when undefined and no prefixes in snippet', + configPrefixes: undefined, + input: 'x = 1', + expected: null, + }, + { + label: 'calls with empty set for full file runs even when no prefixes found', + configPrefixes: undefined, + input: mockTextDoc('x = 1'), + expected: new Set(), + }, + { + label: 'setting takes precedence over meta_import in code', + configPrefixes: ['forced'], + input: 'deephaven_enterprise.controller_import.meta_import("otherPrefix")\n', + expected: new Set(['forced']), + }, + { + label: 'setting takes precedence over meta_import in text doc', + configPrefixes: ['forced'], + input: mockTextDoc( + 'deephaven_enterprise.controller_import.meta_import("otherPrefix")\n' + ), + expected: new Set(['forced']), + }, + ])('$label', async ({ configPrefixes, input, expected }) => { + vi.mocked(ConfigService.getImportPrefixes).mockReturnValue(configPrefixes); + + const service = createTestDhcService({ + setControllerImportPrefixes: mockSetControllerImportPrefixes, + }); + + await service.runCode(input, 'python'); + + if (expected == null) { + expect(mockSetControllerImportPrefixes).not.toHaveBeenCalled(); + } else { + expect(mockSetControllerImportPrefixes).toHaveBeenCalledWith(expected); + } + }); +}); diff --git a/src/services/DhcService.ts b/src/services/DhcService.ts index 8bbc88694..c679b47b2 100644 --- a/src/services/DhcService.ts +++ b/src/services/DhcService.ts @@ -2,6 +2,7 @@ import * as vscode from 'vscode'; import { isAggregateError } from '@deephaven/jsapi-nodejs'; import type { dh as DhcType } from '@deephaven/jsapi-types'; import { + extractControllerImportPrefixes, formatTimestamp, getCombinedRangeLinesText, isNonEmptyArray, @@ -42,6 +43,7 @@ import { } from '../dh/errorUtils'; import { hasErrorCode } from '../util/typeUtils'; import { DisposableBase } from './DisposableBase'; +import { ConfigService } from './ConfigService'; import { assertDefined } from '../shared'; import type { RemoteFileSourceService } from './RemoteFileSourceService'; @@ -522,6 +524,24 @@ export class DhcService extends DisposableBase implements IDhcService { this.isRunningCode = true; if (this.pythonRemoteFileSourcePlugin != null) { + const isDoc = typeof documentOrText !== 'string'; + + // Check for setting override first + const configPrefixes = ConfigService.getImportPrefixes(); + const controllerImportPrefixes = + configPrefixes != null + ? new Set(configPrefixes) + : extractControllerImportPrefixes(text); + + // Update prefixes if: + // 1. Running full file, OR + // 2. Setting or extracted prefixes exist + if (isDoc || controllerImportPrefixes.size > 0) { + this.remoteFileSourceService.setControllerImportPrefixes( + controllerImportPrefixes + ); + } + await this.remoteFileSourceService.setPythonServerExecutionContext( this.cnId, this.session diff --git a/src/services/FilteredWorkspace.ts b/src/services/FilteredWorkspace.ts index 45997c92a..441c40dbe 100644 --- a/src/services/FilteredWorkspace.ts +++ b/src/services/FilteredWorkspace.ts @@ -565,4 +565,9 @@ export class FilteredWorkspace< childMap.set(node.uri, node as FilteredWorkspaceNode); } } + + protected override async onDisposing(): Promise { + this._onDidChangeFileDecorations.dispose(); + this._onDidUpdate.dispose(); + } } diff --git a/src/services/RemoteFileSourceService.spec.ts b/src/services/RemoteFileSourceService.spec.ts new file mode 100644 index 000000000..448b37d6f --- /dev/null +++ b/src/services/RemoteFileSourceService.spec.ts @@ -0,0 +1,338 @@ +/* eslint-disable @typescript-eslint/naming-convention */ +import * as vscode from 'vscode'; +import { beforeEach, describe, it, expect, vi } from 'vitest'; +import { RemoteFileSourceService } from './RemoteFileSourceService'; +import { getClearControllerPrefixesScript } from '../util'; +import type { + FilteredWorkspace, + FilteredWorkspaceTopLevelMarkedNode, +} from './FilteredWorkspace'; +import type { dh as DhcType } from '@deephaven/jsapi-types'; +import type { GroovyPackageName, PythonModuleFullname } from '../types'; + +vi.mock('vscode'); +vi.mock('../util', async () => { + const actual = await vi.importActual('../util'); + return { ...actual, getClearControllerPrefixesScript: vi.fn().mockReturnValue('') }; +}); + +// ── setup ──────────────────────────────────────────────────────────────────── + +const groovyWorkspace = { + onDidUpdate: vi.fn().mockReturnValue({ dispose: vi.fn() }), +} as unknown as FilteredWorkspace; + +const pythonWorkspace = { + onDidUpdate: vi.fn().mockReturnValue({ dispose: vi.fn() }), + getTopLevelMarkedFolders: vi.fn().mockReturnValue([]), + hasFolder: vi.fn(), + hasFile: vi.fn(), +} as unknown as FilteredWorkspace; + +const controllerPrefix = 'controller'; +const customPrefix = 'custom'; + +const myModuleName = 'mymodule' as PythonModuleFullname; +const otherModuleName = 'othermodule' as PythonModuleFullname; +const myModuleFolder = topLevelMarkedFolder(myModuleName); +const otherModuleFolder = topLevelMarkedFolder(otherModuleName); + +function topLevelMarkedFolder( + name: PythonModuleFullname +): FilteredWorkspaceTopLevelMarkedNode { + return { + name, + type: 'topLevelMarkedFolder', + languageId: 'python', + isMarked: true, + uri: vscode.Uri.parse(`file:///path/to/${name}`), + } as const; +} + +beforeEach(() => { + vi.clearAllMocks(); +}); + +// ── tests ──────────────────────────────────────────────────────────────────── + +describe('RemoteFileSourceService', () => { + it('constructor subscribes to workspace updates', () => { + const service = new RemoteFileSourceService( + groovyWorkspace, + pythonWorkspace + ); + expect(service).toBeInstanceOf(RemoteFileSourceService); + expect(vi.mocked(groovyWorkspace.onDidUpdate)).toHaveBeenCalled(); + expect(vi.mocked(pythonWorkspace.onDidUpdate)).toHaveBeenCalled(); + }); + + describe('getPythonTopLevelModuleNames', () => { + it.each([ + { + label: 'no prefixes configured', + folders: [myModuleFolder, otherModuleFolder], + prefixes: [], + }, + { + label: 'default "controller" prefix', + folders: [myModuleFolder, otherModuleFolder], + prefixes: [controllerPrefix], + }, + { + label: 'custom prefix', + folders: [myModuleFolder], + prefixes: [customPrefix], + }, + { + label: 'multiple prefixes', + folders: [myModuleFolder], + prefixes: [controllerPrefix, customPrefix], + }, + { + label: 'no folders marked', + folders: [], + prefixes: [controllerPrefix], + }, + ])('returns $label', ({ folders, prefixes }) => { + vi.mocked(pythonWorkspace.getTopLevelMarkedFolders).mockReturnValue( + folders + ); + + const service = new RemoteFileSourceService( + groovyWorkspace, + pythonWorkspace + ); + + service.setControllerImportPrefixes(new Set(prefixes)); + + const result = service.getPythonTopLevelModuleNames(); + + const expected = new Set([ + ...folders.map(folder => folder.name), + ...prefixes.flatMap(prefix => + folders.map(folder => `${prefix}.${folder.name}`) + ), + ]); + + expect(result).toEqual(expected); + }); + }); + + describe('getPythonModuleSpecData', () => { + it.each([ + { + label: 'regular module found', + moduleFullname: 'mymodule.submodule', + modulePath: 'submodule', + hasFileResult: true, + }, + { + label: 'nested module found', + moduleFullname: 'mymodule.sub.nested', + modulePath: 'sub/nested', + hasFileResult: true, + }, + { + label: 'regular package found', + moduleFullname: 'mymodule.subpackage', + modulePath: 'subpackage', + hasFolderResult: true, + hasFileResult: true, // __init__.py exists + }, + { + label: 'namespace package found', + moduleFullname: 'mymodule.namespace', + modulePath: 'namespace', + hasFolderResult: true, + }, + { + label: 'module with controller prefix stripped', + moduleFullname: 'controller.mymodule.test', + modulePath: 'test', + prefixes: [controllerPrefix], + hasFileResult: true, + }, + { + label: 'bare prefix with no module after (not found)', + moduleFullname: controllerPrefix, + prefixes: [controllerPrefix], + }, + { + label: 'top-level package', + moduleFullname: myModuleName, + modulePath: '', + hasFolderResult: true, + hasFileResult: true, + }, + { + label: 'top-level package with same name as prefix', + moduleFullname: myModuleName, + modulePath: '', + prefixes: [myModuleName], + hasFolderResult: true, + hasFileResult: true, + }, + { + label: 'top-level folder not found', + moduleFullname: 'notfound.module', + }, + { + label: 'module not found', + moduleFullname: 'mymodule.notfound', + }, + ])( + 'returns $label', + ({ + moduleFullname, + modulePath, + prefixes = [], + hasFolderResult = false, + hasFileResult = false, + }) => { + vi.mocked(pythonWorkspace.getTopLevelMarkedFolders).mockReturnValue([ + myModuleFolder, + ]); + vi.mocked(pythonWorkspace.hasFolder).mockReturnValue(hasFolderResult); + vi.mocked(pythonWorkspace.hasFile).mockReturnValue(hasFileResult); + + const service = new RemoteFileSourceService( + groovyWorkspace, + pythonWorkspace + ); + + if (prefixes.length > 0) { + service.setControllerImportPrefixes(new Set(prefixes)); + } + + const result = service.getPythonModuleSpecData( + moduleFullname as PythonModuleFullname + ); + + // Derive expected result from inputs + let expected = null; + if (hasFolderResult || hasFileResult) { + const folderPath = myModuleFolder.uri.fsPath; + + if (!hasFolderResult && hasFileResult) { + // Regular module + expected = { + name: moduleFullname, + isPackage: false, + origin: `${folderPath}/${modulePath}.py`, + }; + } else if (hasFolderResult && hasFileResult) { + // Regular package (with __init__.py) + const packagePath = + modulePath === '' ? folderPath : `${folderPath}/${modulePath}`; + expected = { + name: moduleFullname, + isPackage: true, + origin: `${packagePath}/__init__.py`, + subModuleSearchLocations: [packagePath], + }; + } else if (hasFolderResult && !hasFileResult) { + // Namespace package (without __init__.py) + const packagePath = + modulePath === '' ? folderPath : `${folderPath}/${modulePath}`; + expected = { + name: moduleFullname, + isPackage: true, + origin: undefined, + subModuleSearchLocations: [packagePath], + }; + } + } + + expect(result).toEqual(expected); + } + ); + }); + + it('fires onDidUpdatePythonModuleMeta event when python workspace updates', () => { + const service = new RemoteFileSourceService( + groovyWorkspace, + pythonWorkspace + ); + + const listener = vi.fn(); + service.onDidUpdatePythonModuleMeta(listener); + + const updateListener = vi.mocked(pythonWorkspace.onDidUpdate).mock + .calls[0][0]; + updateListener(); + + expect(listener).toHaveBeenCalled(); + }); +}); + +describe('setControllerImportPrefixes eviction', () => { + const mockSession = { + runCode: vi.fn().mockResolvedValue(undefined), + } as unknown as DhcType.IdeSession; + + // getClearControllerPrefixesScript receives the live Set reference, which is + // cleared in-place immediately after the call. Snapshot copies are captured + // here so assertions see the value at call time, not after the clear. + const capturedEvictions: Set[] = []; + + beforeEach(() => { + capturedEvictions.length = 0; + vi.mocked(getClearControllerPrefixesScript).mockImplementation(set => { + capturedEvictions.push(new Set(set)); + return ''; + }); + }); + + it.each([ + { + label: 'initial single prefix', + prefixCalls: [new Set(['controller'])], + expectedEvicted: new Set(['controller']), + }, + { + label: 'same prefix set twice', + prefixCalls: [new Set(['controller']), new Set(['controller'])], + expectedEvicted: new Set(['controller']), + }, + { + label: 'prefix change unions old and new', + prefixCalls: [new Set(['controller']), new Set(['custom'])], + expectedEvicted: new Set(['controller', 'custom']), + }, + { + label: 'empty prefix set', + prefixCalls: [new Set()], + expectedEvicted: new Set(), + }, + ])( + 'passes accumulated prefixes to eviction script: $label', + async ({ prefixCalls, expectedEvicted }) => { + const service = new RemoteFileSourceService( + groovyWorkspace, + pythonWorkspace + ); + + for (const prefixes of prefixCalls) { + service.setControllerImportPrefixes(prefixes); + } + + await service.setPythonServerExecutionContext(null, mockSession); + + expect(capturedEvictions[0]).toEqual(expectedEvicted); + } + ); + + it('clears eviction set after setPythonServerExecutionContext', async () => { + const service = new RemoteFileSourceService( + groovyWorkspace, + pythonWorkspace + ); + + service.setControllerImportPrefixes(new Set(['controller'])); + await service.setPythonServerExecutionContext(null, mockSession); + await service.setPythonServerExecutionContext(null, mockSession); + + expect(capturedEvictions[0]).toEqual(new Set(['controller'])); + expect(capturedEvictions[1]).toEqual(new Set()); + }); +}); diff --git a/src/services/RemoteFileSourceService.ts b/src/services/RemoteFileSourceService.ts index 54c2c1bdd..6c1072981 100644 --- a/src/services/RemoteFileSourceService.ts +++ b/src/services/RemoteFileSourceService.ts @@ -8,6 +8,7 @@ import { registerGroovyRemoteFileSourcePluginMessageListener, registerPythonRemoteFileSourcePluginMessageListener, getGroovyTopLevelPackageName, + getClearControllerPrefixesScript, } from '../util'; import type { GroovyPackageName, @@ -42,6 +43,8 @@ export class RemoteFileSourceService extends DisposableBase { } private _isGroovyWorkspaceDirty = false; + private _controllerImportPrefixes = new Set(); + private _evictControllerImportPrefixes = new Set(); private _onDidUpdatePythonModuleMeta = new vscode.EventEmitter(); readonly onDidUpdatePythonModuleMeta = @@ -92,7 +95,16 @@ export class RemoteFileSourceService extends DisposableBase { getPythonModuleSpecData( moduleFullname: PythonModuleFullname ): PythonModuleSpecData | null { - const [firstModuleToken, ...restModuleTokens] = moduleFullname.split('.'); + let [firstModuleToken, ...restModuleTokens] = moduleFullname.split('.'); + + // Check if first token is a controller import prefix and strip it + if ( + this._controllerImportPrefixes.has(firstModuleToken) && + restModuleTokens.length > 0 + ) { + firstModuleToken = restModuleTokens[0]; + restModuleTokens = restModuleTokens.slice(1); + } // Get the top-level folder URI that could contain this module const topLevelFolderUri = this._pythonWorkspace @@ -163,7 +175,13 @@ export class RemoteFileSourceService extends DisposableBase { const set = new Set(); this._pythonWorkspace.getTopLevelMarkedFolders().forEach(({ uri }) => { - set.add(getPythonTopLevelModuleFullname(uri)); + const moduleName = getPythonTopLevelModuleFullname(uri); + + set.add(moduleName); + + for (const prefix of this._controllerImportPrefixes) { + set.add(`${prefix}.${moduleName}` as PythonModuleFullname); + } }); return set; @@ -226,6 +244,29 @@ export class RemoteFileSourceService extends DisposableBase { }; } + /** + * Update controller import prefixes based on Python code being executed. + * @param controllerImportPrefixes The set of controller import prefixes to + * use for resolving imports in the code being executed. + */ + setControllerImportPrefixes(controllerImportPrefixes: Set): void { + // Mark previous and current prefixes to evict in the next execution context + // update. + [...this._controllerImportPrefixes, ...controllerImportPrefixes].forEach( + prefix => { + this._evictControllerImportPrefixes.add(prefix); + } + ); + + if (this._evictControllerImportPrefixes.size > 0) { + logger.debug(`marking controller import prefixes for eviction:`, [ + ...this._evictControllerImportPrefixes, + ]); + } + + this._controllerImportPrefixes = controllerImportPrefixes; + } + /** * Set the Groovy server execution context for the plugin. * @param pluginService The remote file source plugin service. @@ -250,8 +291,14 @@ export class RemoteFileSourceService extends DisposableBase { await pluginService.setExecutionContext(isDirty, resourcePaths); } + private _pythonSetExecutionContextI = 0; + private _pythonExecutionContextQueue: Promise = Promise.resolve(); + /** * Set the Python server execution context for the plugin using the given session. + * We use a Promise queue to ensure that execution context updates are processed + * sequentially. This is mostly to prevent Python workspace events that call + * this method without awaiting the response from clearing the execution context. * @param connectionId The unique ID of the connection. * @param session The IdeSession to use to run the code. */ @@ -259,11 +306,37 @@ export class RemoteFileSourceService extends DisposableBase { connectionId: UniqueID | null, session: DhcType.IdeSession ): Promise { - const setExecutionContextScript = getSetExecutionContextScript( - connectionId, - this.getPythonTopLevelModuleNames() - ); + const label = `setPythonServerExecutionContext: ${++this._pythonSetExecutionContextI}:${connectionId}`; + + logger.debug(`${label}: queuing`); + + this._pythonExecutionContextQueue = this._pythonExecutionContextQueue + // Ignore errors from previous calls. They will get raised to the caller + // that queued them, but we dont' want them to break the chain + .catch(() => {}) + .then(async () => { + logger.debug(`${label}: running`); + + const clearControllerPrefixesScript = getClearControllerPrefixesScript( + this._evictControllerImportPrefixes + ); + this._evictControllerImportPrefixes.clear(); + + const setExecutionContextScript = getSetExecutionContextScript( + connectionId, + this.getPythonTopLevelModuleNames() + ); + + const scripts = [ + clearControllerPrefixesScript, + setExecutionContextScript, + ].filter(Boolean); + + await session.runCode(scripts.join('\n')); + }); + + await this._pythonExecutionContextQueue; - await session.runCode(setExecutionContextScript); + logger.debug(`${label}: complete`); } } diff --git a/src/types/serviceTypes.d.ts b/src/types/serviceTypes.d.ts index 919f3db08..3e7f2890b 100644 --- a/src/types/serviceTypes.d.ts +++ b/src/types/serviceTypes.d.ts @@ -44,6 +44,7 @@ export interface IConfigService { isMcpEnabled: () => boolean; getCoreServers: () => CoreConnectionConfig[]; getEnterpriseServers: () => EnterpriseConnectionConfig[]; + getImportPrefixes: () => string[] | undefined; getMcpAutoUpdateConfig: () => boolean; setMcpAutoUpdateConfig: (value: boolean) => Promise; toggleMcp: (enable?: boolean) => Promise; diff --git a/src/util/__snapshots__/remoteFileSourceUtils.spec.ts.snap b/src/util/__snapshots__/remoteFileSourceUtils.spec.ts.snap index a673f0f4c..36e031e68 100644 --- a/src/util/__snapshots__/remoteFileSourceUtils.spec.ts.snap +++ b/src/util/__snapshots__/remoteFileSourceUtils.spec.ts.snap @@ -1,5 +1,19 @@ // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html +exports[`getClearControllerPrefixesScript > empty set 1`] = `""`; + +exports[`getClearControllerPrefixesScript > multiple prefixes 1`] = ` +"modules_to_delete = [name for name in sys.modules if name == "a" or name.startswith("a.") or name == "b" or name.startswith("b.")] +for module in modules_to_delete: + del sys.modules[module]" +`; + +exports[`getClearControllerPrefixesScript > single prefix 1`] = ` +"modules_to_delete = [name for name in sys.modules if name == "controller" or name.startswith("controller.")] +for module in modules_to_delete: + del sys.modules[module]" +`; + exports[`getFileTreeItem > should return a TreeItem for a file element 1`] = ` { "collapsibleState": 0, diff --git a/src/util/index.ts b/src/util/index.ts index 6a5951058..29043971e 100644 --- a/src/util/index.ts +++ b/src/util/index.ts @@ -13,6 +13,7 @@ export * from './maps'; export * from './OutputChannelWithHistory'; export * from './panelUtils'; export * from './promiseUtils'; +export * from './pythonUtils'; export * from './remoteFileSourceMsgUtils'; export * from './remoteFileSourceUtils'; export * from './sanitizeUtils'; diff --git a/src/util/pythonUtils.spec.ts b/src/util/pythonUtils.spec.ts new file mode 100644 index 000000000..f74536511 --- /dev/null +++ b/src/util/pythonUtils.spec.ts @@ -0,0 +1,57 @@ +import { describe, it, expect } from 'vitest'; +import { extractControllerImportPrefixes } from './pythonUtils'; + +// Code generators for each pattern's call syntax +const makeCode = { + pattern1: (args: string): string => + `deephaven_enterprise.controller_import.meta_import(${args})`, + pattern2: (args: string): string => + `from deephaven_enterprise import controller_import\ncontroller_import.meta_import(${args})`, + pattern3: (args: string): string => + `from deephaven_enterprise.controller_import import meta_import\nmeta_import(${args})`, +}; + +describe('extractControllerImportPrefixes', () => { + // Shared behavior across all 3 patterns + describe.each([ + ['pattern 1 (direct call)', makeCode.pattern1], + ['pattern 2 (from module import)', makeCode.pattern2], + ['pattern 3 (from function import)', makeCode.pattern3], + ])('%s', (_, make: (args: string) => string) => { + it.each([ + ['double-quoted prefix', '"my_prefix"', new Set(['my_prefix'])], + ["single-quoted prefix", `'my_prefix'`, new Set(['my_prefix'])], + ['no arg defaults to "controller"', '', new Set(['controller'])], + ['whitespace inside parens', ' "spaced" ', new Set(['spaced'])], + ])('%s', (_label, args, expected) => { + expect(extractControllerImportPrefixes(make(args))).toEqual(expected); + }); + + it('multiple calls yield multiple prefixes', () => { + expect( + extractControllerImportPrefixes(`${make('"a"')}\n${make('"b"')}`) + ).toEqual(new Set(['a', 'b'])); + }); + }); + + // Unique to patterns 2 & 3: call without import statement yields no match + it.each([ + ['pattern 2', `controller_import.meta_import("foo")`], + ['pattern 3', `meta_import("foo")`], + ])('%s - call without import statement yields empty set', (_, code) => { + expect(extractControllerImportPrefixes(code)).toEqual(new Set()); + }); + + // Edge cases + it.each([ + [ + 'duplicate prefix across patterns is deduplicated', + `from deephaven_enterprise.controller_import import meta_import\ndeephaven_enterprise.controller_import.meta_import("dup")\nmeta_import("dup")`, + new Set(['dup']), + ], + ['no meta_import calls', 'x = 1', new Set()], + ['empty string', '', new Set()], + ])('%s', (_label, code, expected) => { + expect(extractControllerImportPrefixes(code)).toEqual(expected); + }); +}); diff --git a/src/util/pythonUtils.ts b/src/util/pythonUtils.ts new file mode 100644 index 000000000..1091529d0 --- /dev/null +++ b/src/util/pythonUtils.ts @@ -0,0 +1,59 @@ +const DEFAULT_META_IMPORT_PREFIX = 'controller' as const; + +/** + * Extract controller prefixes from Python source code that uses + * `deephaven_enterprise.controller_import.meta_import()`. + * + * Supported patterns: + * 1. `import deephaven_enterprise.controller_import` + `deephaven_enterprise.controller_import.meta_import()` call + * 2. `from deephaven_enterprise import controller_import` + `controller_import.meta_import()` call + * 3. `from deephaven_enterprise.controller_import import meta_import` + `meta_import()` call + * + * Limitations: + * - Import aliases are not detected + * - Multiline calls are not detected + * + * @param pythonCode The Python source code to scan. + * @returns A set of controller import prefixes found in the code. + */ +export function extractControllerImportPrefixes( + pythonCode: string +): Set { + const prefixes = new Set(); + + // Pattern 1: deephaven_enterprise.controller_import.meta_import() direct call + const directCallPattern = + /deephaven_enterprise\.controller_import\.meta_import\(\s*(?:["'](\w+)["'])?\s*\)/g; + let match: RegExpExecArray | null; + + while ((match = directCallPattern.exec(pythonCode)) !== null) { + prefixes.add(match[1] ?? DEFAULT_META_IMPORT_PREFIX); + } + + // Pattern 2: from deephaven_enterprise import controller_import + // followed by controller_import.meta_import() call + const fromModulePattern = + /from\s+deephaven_enterprise\s+import\s+controller_import/; + if (fromModulePattern.test(pythonCode)) { + const callPattern = + /controller_import\.meta_import\(\s*(?:["'](\w+)["'])?\s*\)/g; + + while ((match = callPattern.exec(pythonCode)) !== null) { + prefixes.add(match[1] ?? DEFAULT_META_IMPORT_PREFIX); + } + } + + // Pattern 3: from deephaven_enterprise.controller_import import meta_import + // followed by meta_import() call + const fromImportPattern = + /from\s+deephaven_enterprise\.controller_import\s+import\s+meta_import/; + if (fromImportPattern.test(pythonCode)) { + const callPattern = /\bmeta_import\(\s*(?:["'](\w+)["'])?\s*\)/g; + + while ((match = callPattern.exec(pythonCode)) !== null) { + prefixes.add(match[1] ?? DEFAULT_META_IMPORT_PREFIX); + } + } + + return prefixes; +} diff --git a/src/util/remoteFileSourceUtils.spec.ts b/src/util/remoteFileSourceUtils.spec.ts index 6d5d475f4..792dade42 100644 --- a/src/util/remoteFileSourceUtils.spec.ts +++ b/src/util/remoteFileSourceUtils.spec.ts @@ -2,6 +2,7 @@ import * as vscode from 'vscode'; import { beforeEach, describe, it, expect, vi } from 'vitest'; import type { dh as DhcType } from '@deephaven/jsapi-types'; import { + getClearControllerPrefixesScript, getFileTreeItem, getFolderTreeItem, getGroovyTopLevelPackageName, @@ -332,3 +333,13 @@ describe('sendWidgetMessageAsync', () => { expect(removeEventListener).toHaveBeenCalled(); }); }); + +describe('getClearControllerPrefixesScript', () => { + it.each([ + ['empty set', new Set()], + ['single prefix', new Set(['controller'])], + ['multiple prefixes', new Set(['a', 'b'])], + ])('%s', (_label, prefixes) => { + expect(getClearControllerPrefixesScript(prefixes)).toMatchSnapshot(); + }); +}); diff --git a/src/util/remoteFileSourceUtils.ts b/src/util/remoteFileSourceUtils.ts index 56738c605..9748af358 100644 --- a/src/util/remoteFileSourceUtils.ts +++ b/src/util/remoteFileSourceUtils.ts @@ -315,6 +315,43 @@ export async function getWorkspaceFileUriMap( return map; } +/** + * There are certain import scenarios where a controller prefix import can get cached + * in a way that doesn't get evicted when its dependencies do. + * + * e.g. + * - Run entry_point_script.py + * → imports prefix.package1 + * → imports prefix.package2 (marked as remote source) + * - Both cached in sys.modules + * - Unmark prefix.package2 as remote source + * - Plugin evicts prefix.package2 from cache, but prefix.package1 is still + * cached and holds stale reference to prefix.package2 + * + * Solution: Clear all `prefix` and `prefix.*` from sys.modules. + * + * @param controllerPrefixes + * @returns A Python script that will clear all modules from sys.modules that + * start with any of the given prefixes. + */ +export function getClearControllerPrefixesScript( + controllerPrefixes: Set +): string { + if (controllerPrefixes.size === 0) { + return ''; + } + + const startsWithConditions = [...controllerPrefixes].map( + prefix => `name == "${prefix}" or name.startswith("${prefix}.")` + ); + + return [ + `modules_to_delete = [name for name in sys.modules if ${startsWithConditions.join(' or ')}]`, + 'for module in modules_to_delete:', + ' del sys.modules[module]', + ].join('\n'); +} + /** * Get a script to set the execution context on the remote file source plugin. * @param connectionId The unique ID of the connection. diff --git a/src/util/sets/SerializedKeySet.ts b/src/util/sets/SerializedKeySet.ts index 1881a9a74..b71e85409 100644 --- a/src/util/sets/SerializedKeySet.ts +++ b/src/util/sets/SerializedKeySet.ts @@ -138,4 +138,11 @@ export abstract class SerializedKeySet implements IDisposable { *values(): IterableIterator { yield* this.keys(); } + + /** + * Default iterator for the set. Returns the same as values(). + */ + [Symbol.iterator](): IterableIterator { + return this.values(); + } }