fix: isInWorkspace should work on closed files. (#1004)

* fix: inWorkspace only worked on open files

* refactor: inject lsp instead of workspaceFolders

* fix: use proper workspace path in tests

* fix: stub lsp initializeParams in exec bash

* fix: isInWorkspace returns true for workspace itself

Author: Hweinstock

core/aws-lsp-core/src/util/workspaceUtils.test.ts

@@ -3,14 +3,57 @@ import * as fs from 'fs/promises'
 import * as assert from 'assert'
 import * as path from 'path'
 import { TestFolder } from '../test/testFolder'
-import { readDirectoryRecursively, getEntryPath } from './workspaceUtils'
+import { readDirectoryRecursively, getEntryPath, isParentFolder, isInWorkspace } from './workspaceUtils'
 import { TestFeatures } from '@aws/language-server-runtimes/testing'
 
 describe('workspaceUtils', function () {
     beforeEach(function () {
         mockfs.restore()
     })
 
+    describe('isParentFolder', function () {
+        it('handles different cases', function () {
+            assert.ok(isParentFolder('/foo', '/foo/bar'), 'simple case')
+            assert.ok(isParentFolder('/foo', '/foo/bar/'), 'trailing slash in child')
+            assert.ok(isParentFolder('/foo', '/foo/bar.txt'), 'files')
+            assert.ok(isParentFolder('/foo', '/foo/bar/baz'), 'neseted directory')
+            assert.ok(!isParentFolder('/foo', '/foo'), 'duplicates')
+            assert.ok(!isParentFolder('/foo', '/foobar'), 'prefixing')
+            assert.ok(!isParentFolder('/foo/bar', '/foo'), 'reverse')
+            assert.ok(isParentFolder('/foo/', '/foo/bar/baz/'), 'trailing slash in both')
+        })
+    })
+
+    describe('isInWorkspace', function () {
+        it('finds the file within the workspace', function () {
+            const workspaceFolders = ['/foo']
+
+            const positiveFilePath = '/foo/bar/baz.txt'
+            const negativeFilePath = '/notfoo/bar/baz.txt'
+
+            assert.ok(isInWorkspace(workspaceFolders, positiveFilePath), 'file is within the workspace')
+            assert.ok(!isInWorkspace(workspaceFolders, negativeFilePath), 'file is not within the workspace')
+        })
+
+        it('handles multi-root workspaces', function () {
+            const workspaceFolders = ['/foo', '/bar']
+
+            const positiveFilePath = '/foo/bar/baz.txt'
+            const secondPositiveFilePath = '/bar/bax.txt'
+            const negativeFilePath = '/notfoo/bar/baz.txt'
+
+            assert.ok(isInWorkspace(workspaceFolders, positiveFilePath), 'file is within the workspace')
+            assert.ok(isInWorkspace(workspaceFolders, secondPositiveFilePath), 'file is within the workspace')
+            assert.ok(!isInWorkspace(workspaceFolders, negativeFilePath), 'file is not within the workspace')
+        })
+
+        it('handles the case where its the workspace itself', function () {
+            const workspaceFolders = ['/foo']
+
+            assert.ok(isInWorkspace(workspaceFolders, '/foo'), 'workspace is inside itself')
+        })
+    })
+
     describe('readDirectoryRecursively', function () {
         let tempFolder: TestFolder
         let testFeatures: TestFeatures

core/aws-lsp-core/src/util/workspaceUtils.ts

@@ -85,10 +85,18 @@ export function getEntryPath(entry: Dirent) {
 }
 
 // TODO: port this to runtimes?
-export function getWorkspaceFolders(lsp: Features['lsp']): string[] {
+export function getWorkspaceFolderPaths(lsp: Features['lsp']): string[] {
     return lsp.getClientInitializeParams()?.workspaceFolders?.map(({ uri }) => URI.parse(uri).fsPath) ?? []
 }
 
-export async function inWorkspace(workspace: Features['workspace'], filepath: string) {
-    return (await workspace.getTextDocument(URI.file(filepath).toString())) !== undefined
+export function isParentFolder(parentPath: string, childPath: string): boolean {
+    const normalizedParentPath = path.normalize(parentPath)
+    const normalizedChildPath = path.normalize(childPath)
+
+    const relative = path.relative(normalizedParentPath, normalizedChildPath)
+    return relative !== '' && !relative.startsWith('..') && !path.isAbsolute(relative)
+}
+
+export function isInWorkspace(workspaceFolderPaths: string[], filepath: string) {
+    return workspaceFolderPaths.some(wsFolder => isParentFolder(wsFolder, filepath) || wsFolder === filepath)
 }

server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/agenticChatTriggerContext.ts

@@ -66,7 +66,7 @@ export class AgenticChatTriggerContext {
         additionalContent?: AdditionalContentEntryAddition[]
     ): GenerateAssistantResponseCommandInput {
         const { prompt } = params
-        const defaultEditorState = { workspaceFolders: workspaceUtils.getWorkspaceFolders(this.#lsp) }
+        const defaultEditorState = { workspaceFolders: workspaceUtils.getWorkspaceFolderPaths(this.#lsp) }
         const data: GenerateAssistantResponseCommandInput = {
             conversationState: {
                 chatTriggerType: chatTriggerType,

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/executeBash.test.ts

@@ -5,12 +5,18 @@ import { ExecuteBash } from './executeBash'
 import { Features } from '@aws/language-server-runtimes/server-interface/server'
 import { TestFeatures } from '@aws/language-server-runtimes/testing'
 import { TextDocument } from 'vscode-languageserver-textdocument'
+import { URI } from 'vscode-uri'
+import { InitializeParams } from '@aws/language-server-runtimes/protocol'
 
 describe('ExecuteBash Tool', () => {
-    let features: Features
+    let features: TestFeatures
+    const workspaceFolder = '/workspace/folder'
 
     before(function () {
         features = new TestFeatures()
+        features.lsp.getClientInitializeParams.returns({
+            workspaceFolders: [{ uri: URI.file(workspaceFolder).toString(), name: 'test' }],
+        } as InitializeParams)
     })
 
     beforeEach(() => {
@@ -76,7 +82,7 @@ describe('ExecuteBash Tool', () => {
         })
         const result = await execBash.requiresAcceptance({
             command: 'cat /not/in/workspace/file.txt',
-            cwd: '/workspace/folder',
+            cwd: workspaceFolder,
         })
 
         assert.equal(
@@ -95,7 +101,7 @@ describe('ExecuteBash Tool', () => {
                 getTextDocument: async s => ({}) as TextDocument,
             },
         })
-        const result = await execBash.requiresAcceptance({ command: 'cat ./file.txt', cwd: '/workspace/folder' })
+        const result = await execBash.requiresAcceptance({ command: 'cat ./file.txt', cwd: workspaceFolder })
 
         assert.equal(result.requiresAcceptance, false, 'Relative path inside workspace should not require acceptance')
     })
@@ -108,7 +114,7 @@ describe('ExecuteBash Tool', () => {
                 getTextDocument: async s => ({}) as TextDocument,
             },
         })
-        const result = await execBash.requiresAcceptance({ command: 'echo hello world', cwd: '/workspace/folder' })
+        const result = await execBash.requiresAcceptance({ command: 'echo hello world', cwd: workspaceFolder })
 
         assert.equal(
             result.requiresAcceptance,

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/executeBash.ts

@@ -9,6 +9,7 @@ import { ChildProcess, ChildProcessOptions } from '@aws/lsp-core/out/util/proces
 // eslint-disable-next-line import/no-nodejs-modules
 import { isAbsolute, join } from 'path' // Safe to import on web since this is part of path-browserify
 import { Features } from '../../types'
+import { getWorkspaceFolderPaths } from '@aws/lsp-core/out/util/workspaceUtils'
 
 export enum CommandCategory {
     ReadOnly,
@@ -122,11 +123,11 @@ interface TimestampedChunk {
 
 export class ExecuteBash {
     private childProcess?: ChildProcess
-    private readonly workspace: Features['workspace']
     private readonly logging: Features['logging']
-    constructor(features: Pick<Features, 'logging' | 'workspace'> & Partial<Features>) {
-        this.workspace = features.workspace
+    private readonly lsp: Features['lsp']
+    constructor(features: Pick<Features, 'logging' | 'lsp'> & Partial<Features>) {
         this.logging = features.logging
+        this.lsp = features.lsp
     }
 
     public async validate(command: string): Promise<void> {
@@ -195,7 +196,7 @@ export class ExecuteBash {
                     if (this.looksLikePath(arg)) {
                         // If not absolute, resolve using workingDirectory if available.
                         const fullPath = !isAbsolute(arg) && params.cwd ? join(params.cwd, arg) : arg
-                        const isInWorkspace = await workspaceUtils.inWorkspace(this.workspace, fullPath)
+                        const isInWorkspace = workspaceUtils.isInWorkspace(getWorkspaceFolderPaths(this.lsp), fullPath)
                         if (!isInWorkspace) {
                             return { requiresAcceptance: true, warning: destructiveCommandWarningMessage }
                         }

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/listDirectory.ts

@@ -4,6 +4,7 @@ import { workspaceUtils } from '@aws/lsp-core'
 import { Features } from '@aws/language-server-runtimes/server-interface/server'
 import { sanitize } from '@aws/lsp-core/out/util/path'
 import { DEFAULT_EXCLUDE_PATTERNS } from '../../chat/constants'
+import { getWorkspaceFolderPaths } from '@aws/lsp-core/out/util/workspaceUtils'
 
 export interface ListDirectoryParams {
     path: string
@@ -13,10 +14,12 @@ export interface ListDirectoryParams {
 export class ListDirectory {
     private readonly logging: Features['logging']
     private readonly workspace: Features['workspace']
+    private readonly lsp: Features['lsp']
 
-    constructor(features: Pick<Features, 'logging' | 'workspace'>) {
+    constructor(features: Pick<Features, 'logging' | 'workspace' | 'lsp'>) {
         this.logging = features.logging
         this.workspace = features.workspace
+        this.lsp = features.lsp
     }
 
     public async validate(params: ListDirectoryParams): Promise<void> {
@@ -49,7 +52,7 @@ export class ListDirectory {
     }
 
     public async requiresAcceptance(path: string): Promise<CommandValidation> {
-        return { requiresAcceptance: !(await workspaceUtils.inWorkspace(this.workspace, path)) }
+        return { requiresAcceptance: !workspaceUtils.isInWorkspace(getWorkspaceFolderPaths(this.lsp), path) }
     }
 
     public async invoke(params: ListDirectoryParams): Promise<InvokeOutput> {

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolServer.ts

@@ -7,11 +7,11 @@ import { LspGetDocuments, LspGetDocumentsParams } from './lspGetDocuments'
 import { LspReadDocumentContents, LspReadDocumentContentsParams } from './lspReadDocumentContents'
 import { LspApplyWorkspaceEdit, LspApplyWorkspaceEditParams } from './lspApplyWorkspaceEdit'
 
-export const FsToolsServer: Server = ({ workspace, logging, agent }) => {
+export const FsToolsServer: Server = ({ workspace, logging, agent, lsp }) => {
     const fsReadTool = new FsRead({ workspace, logging })
     const fsWriteTool = new FsWrite({ workspace, logging })
 
-    const listDirectoryTool = new ListDirectory({ workspace, logging })
+    const listDirectoryTool = new ListDirectory({ workspace, logging, lsp })
 
     agent.addTool(fsReadTool.getSpec(), async (input: FsReadParams) => {
         // TODO: fill in logic for handling invalid tool invocations
@@ -32,8 +32,8 @@ export const FsToolsServer: Server = ({ workspace, logging, agent }) => {
     return () => {}
 }
 
-export const BashToolsServer: Server = ({ logging, workspace, agent }) => {
-    const bashTool = new ExecuteBash({ logging, workspace })
+export const BashToolsServer: Server = ({ logging, workspace, agent, lsp }) => {
+    const bashTool = new ExecuteBash({ logging, workspace, lsp })
     agent.addTool(bashTool.getSpec(), (input: ExecuteBashParams) => bashTool.invoke(input))
     return () => {}
 }