Expose open LSP documents via $psEditor.Workspace.Documents

src/PowerShellEditorServices/Extensions/EditorWorkspace.cs

@@ -28,6 +28,11 @@ public sealed class EditorWorkspace
         /// </summary>
         public string[] Paths => editorOperations.GetWorkspacePaths();
 
+        /// <summary>
+        /// Get all currently open documents in the workspace.
+        /// </summary>
+        public WorkspaceOpenDocument[] Documents => editorOperations.GetWorkspaceOpenDocuments();
+
         #endregion
 
         #region Constructors
@@ -76,13 +81,15 @@ public sealed class EditorWorkspace
         /// <param name="filePath">The path to the file to be closed.</param>
         [System.Diagnostics.CodeAnalysis.SuppressMessage("Usage", "VSTHRD002:Avoid problematic synchronous waits", Justification = "Supporting synchronous API.")]
         public void CloseFile(string filePath) => editorOperations.CloseFileAsync(filePath).Wait();
+        public void CloseFile(WorkspaceOpenDocument document) => CloseFile(document.Path);
 
         /// <summary>
         /// Saves an open file in the workspace.
         /// </summary>
         /// <param name="filePath">The path to the file to be saved.</param>
         [System.Diagnostics.CodeAnalysis.SuppressMessage("Usage", "VSTHRD002:Avoid problematic synchronous waits", Justification = "Supporting synchronous API.")]
         public void SaveFile(string filePath) => editorOperations.SaveFileAsync(filePath).Wait();
+        public void SaveFile(WorkspaceOpenDocument document) => SaveFile(document.Path);
 
         /// <summary>
         /// Saves a file with a new name AKA a copy.

src/PowerShellEditorServices/Extensions/IEditorOperations.cs

@@ -3,9 +3,34 @@
 
 using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Services.TextDocument;
+#nullable enable
 
 namespace Microsoft.PowerShell.EditorServices.Extensions
 {
+    public readonly struct WorkspaceOpenDocument(string path, bool saved)
+    {
+        /// <summary>
+        /// Gets the path or URI of the open document.
+        /// </summary>
+        public string Path { get; } = path;
+
+        /// <summary>
+        /// Gets whether the document is backed by a saved file path (not in-memory).
+        /// </summary>
+        public bool Saved { get; } = saved;
+
+        /// <summary>
+        /// Gets the display name of this document and unsaved status.
+        /// </summary>
+        /// <returns>The display name of this document.</returns>
+        public override string ToString()
+        {
+            string documentPath = Path ?? string.Empty;
+            string fileName = System.IO.Path.GetFileName(documentPath);
+            return Saved ? fileName : fileName + " [Unsaved]";
+        }
+    }
+
     /// <summary>
     /// Provides an interface that must be implemented by an editor
     /// host to perform operations invoked by extensions written in
@@ -32,6 +57,12 @@ internal interface IEditorOperations
         /// <returns></returns>
         string[] GetWorkspacePaths();
 
+        /// <summary>
+        /// Get all open documents in the current workspace session.
+        /// </summary>
+        /// <returns>All currently open documents.</returns>
+        WorkspaceOpenDocument[] GetWorkspaceOpenDocuments();
+
         /// <summary>
         /// Resolves the given file path relative to the current workspace path.
         /// </summary>

src/PowerShellEditorServices/Services/Extension/EditorOperationsService.cs

@@ -198,6 +198,14 @@ public async Task SaveFileAsync(string currentPath, string newSavePath)
 
         public string[] GetWorkspacePaths() => _workspaceService.WorkspacePaths.ToArray();
 
+        public WorkspaceOpenDocument[] GetWorkspaceOpenDocuments()
+            => [..
+                    _workspaceService
+                    .GetOpenedFiles()
+                    .Where(static scriptFile => scriptFile.IsOpen)
+                    .Select(static scriptFile => new WorkspaceOpenDocument(scriptFile.FilePath, !scriptFile.IsInMemory))
+                ];
+
         public string GetWorkspaceRelativePath(ScriptFile scriptFile) => _workspaceService.GetRelativePath(scriptFile);
 
         public async Task ShowInformationMessageAsync(string message)

src/PowerShellEditorServices/Services/TextDocument/Handlers/CompletionHandler.cs

@@ -366,7 +366,7 @@ private CompletionItem CreateProviderItemCompletion(
             if (textToBeReplaced.IndexOf(PSScriptRootVariable, StringComparison.OrdinalIgnoreCase) is int variableIndex and not -1
                 && System.IO.Path.GetDirectoryName(scriptFile.FilePath) is string scriptFolder and not ""
                 && completionText.IndexOf(scriptFolder, StringComparison.OrdinalIgnoreCase) is int pathIndex and not -1
-                && !scriptFile.IsInMemory)
+                && !scriptFile.IsUntitled)
             {
                 completionText = completionText
                     .Remove(pathIndex, scriptFolder.Length)

src/PowerShellEditorServices/Services/TextDocument/Handlers/TextDocumentHandler.cs

@@ -111,10 +111,11 @@ public override Task<Unit> Handle(DidCloseTextDocumentParams notification, Cance
             {
                 fileToClose.IsOpen = false;
 
-                // If the file watcher is supported, only close in-memory files when this
+                // If the file watcher is supported, only close non-file-backed documents when this
                 // notification is triggered. This lets us keep workspace files open so we can scan
                 // for references. When a file is deleted, the file watcher will close the file.
-                if (!_isFileWatcherSupported || fileToClose.IsInMemory)
+                bool isBackedByFile = !fileToClose.IsUntitled;
+                if (!_isFileWatcherSupported || !isBackedByFile)
                 {
                     _workspaceService.CloseFile(fileToClose);
                 }
@@ -132,6 +133,9 @@ public override async Task<Unit> Handle(DidSaveTextDocumentParams notification,
 
             if (savedFile != null)
             {
+                // On a save, untitled files will remain in memory, so this won't change for those
+                savedFile.IsInMemory = savedFile.IsUntitled;
+
                 if (_remoteFileManagerService.IsUnderRemoteTempPath(savedFile.FilePath))
                 {
                     await _remoteFileManagerService.SaveRemoteFileAsync(savedFile.FilePath).ConfigureAwait(false);

src/PowerShellEditorServices/Services/TextDocument/ScriptFile.cs

@@ -51,9 +51,15 @@ internal sealed class ScriptFile
 
         /// <summary>
         /// Gets a boolean that determines whether this file is
-        /// in-memory or not (either unsaved or non-file content).
+        /// in-memory or not (either unsaved or non-file content) aka "dirty"
         /// </summary>
-        public bool IsInMemory { get; }
+        public bool IsInMemory { get; internal set; }
+
+        /// <summary>
+        /// Gets a value indicating whether the document URI is not a <c>file://</c> URI
+        /// (for example, an <c>untitled:</c> URI).
+        /// </summary>
+        public bool IsUntitled => !DocumentUri.ToUri().IsFile;
 
         /// <summary>
         /// Gets a string containing the full contents of the file.
@@ -105,6 +111,9 @@ public Token[] ScriptTokens
 
         internal ReferenceTable References { get; }
 
+        /// <summary>
+        /// Indicates whether the file is currently open in the editor. PSES may open files for analysis that aren't actually visible in the editor.
+        /// </summary>
         internal bool IsOpen { get; set; }
 
         #endregion
@@ -127,11 +136,15 @@ internal ScriptFile(
             // so that other operations know it's untitled/in-memory
             // and don't think that it's a relative path
             // on the file system.
-            IsInMemory = !docUri.ToUri().IsFile;
+            DocumentUri = docUri;
+
+            // Initial state of document. Untitled files are in memory by definition, otherwise files start non-dirty on a filesystem
+            IsInMemory = IsUntitled;
+
             FilePath = IsInMemory
                 ? docUri.ToString()
                 : docUri.GetFileSystemPath();
-            DocumentUri = docUri;
+
             IsAnalysisEnabled = true;
             this.powerShellVersion = powerShellVersion;
 
@@ -365,6 +378,9 @@ public void ApplyChange(FileChange fileChange)
             // Parse the script again to be up-to-date
             ParseFileContents();
             References.TagAsChanged();
+
+            // Flag the script as modified
+            IsInMemory = true;
         }
 
         /// <summary>

src/PowerShellEditorServices/Services/Workspace/WorkspaceService.cs

@@ -302,7 +302,14 @@ public void CloseFile(ScriptFile scriptFile)
             Validate.IsNotNull(nameof(scriptFile), scriptFile);
 
             string keyName = GetFileKey(scriptFile.DocumentUri);
-            workspaceFiles.TryRemove(keyName, out ScriptFile _);
+            if (workspaceFiles.TryRemove(keyName, out ScriptFile _))
+            {
+                logger.LogDebug("Closed file: " + scriptFile.DocumentUri);
+            }
+            else
+            {
+                logger.LogWarning("Tried to close file that was not open: " + scriptFile.DocumentUri);
+            }
         }
 
         /// <summary>
@@ -312,7 +319,7 @@ public void CloseFile(ScriptFile scriptFile)
         public string GetRelativePath(ScriptFile scriptFile)
         {
             Uri fileUri = scriptFile.DocumentUri.ToUri();
-            if (!scriptFile.IsInMemory)
+            if (!scriptFile.IsUntitled)
             {
                 // Support calculating out-of-workspace relative paths in the common case of a
                 // single workspace folder. Otherwise try to get the matching folder.

test/PowerShellEditorServices.Test/Extensions/EditorOperationsServiceTests.cs

@@ -0,0 +1,98 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System;
+using System.IO;
+using Microsoft.Extensions.Logging.Abstractions;
+using Microsoft.PowerShell.EditorServices.Extensions;
+using Microsoft.PowerShell.EditorServices.Services;
+using Microsoft.PowerShell.EditorServices.Services.Extension;
+using Microsoft.PowerShell.EditorServices.Services.TextDocument;
+using OmniSharp.Extensions.LanguageServer.Protocol;
+using Xunit;
+
+namespace PowerShellEditorServices.Test.Extensions
+{
+    [Trait("Category", "Extensions")]
+    public class EditorOperationsServiceTests
+    {
+        [Fact]
+        public void GetWorkspaceOpenDocumentsReturnsOnlyOpenDocumentsAndCurrentInMemoryState()
+        {
+            WorkspaceService workspaceService = new(NullLoggerFactory.Instance);
+
+            ScriptFile openSaved = CreateFileBuffer(workspaceService, "open-saved.ps1");
+            openSaved.IsOpen = true;
+            openSaved.IsInMemory = false;
+
+            ScriptFile openUnsaved = CreateFileBuffer(workspaceService, "open-unsaved.ps1");
+            openUnsaved.IsOpen = true;
+            openUnsaved.IsInMemory = true;
+
+            ScriptFile closed = CreateFileBuffer(workspaceService, "closed.ps1");
+            closed.IsOpen = false;
+            closed.IsInMemory = false;
+
+            EditorOperationsService editorOperationsService = new(
+                psesHost: null,
+                workspaceService,
+                languageServer: null);
+
+            WorkspaceOpenDocument[] documents = editorOperationsService.GetWorkspaceOpenDocuments();
+
+            Assert.Equal(2, documents.Length);
+            Assert.Contains(documents, static document => document.Path.EndsWith("open-saved.ps1") && document.Saved);
+            Assert.Contains(documents, static document => document.Path.EndsWith("open-unsaved.ps1") && !document.Saved);
+            Assert.DoesNotContain(documents, static document => document.Path.EndsWith("closed.ps1"));
+        }
+
+        [Fact]
+        public void GetWorkspaceOpenDocumentsTracksEditedAndUntitledSaveStates()
+        {
+            WorkspaceService workspaceService = new(NullLoggerFactory.Instance);
+
+            ScriptFile openSaved = CreateFileBuffer(workspaceService, "open-saved.ps1");
+            openSaved.IsOpen = true;
+
+            ScriptFile openUntitled = workspaceService.GetFileBuffer("untitled:Untitled-1", initialBuffer: string.Empty);
+            openUntitled.IsOpen = true;
+
+            EditorOperationsService editorOperationsService = new(
+                psesHost: null,
+                workspaceService,
+                languageServer: null);
+
+            WorkspaceOpenDocument[] initialDocuments = editorOperationsService.GetWorkspaceOpenDocuments();
+            Assert.Contains(initialDocuments, static document => document.Path.EndsWith("open-saved.ps1") && document.Saved);
+            Assert.Contains(initialDocuments, static document => document.Path.StartsWith("untitled:", StringComparison.Ordinal) && !document.Saved);
+
+            openSaved.ApplyChange(new FileChange
+            {
+                Line = 1,
+                Offset = 1,
+                EndLine = 1,
+                EndOffset = 1,
+                InsertString = "Set-StrictMode -Version Latest"
+            });
+            Assert.Contains("Set-StrictMode -Version Latest", openSaved.Contents, StringComparison.Ordinal);
+
+            WorkspaceOpenDocument[] editedDocuments = editorOperationsService.GetWorkspaceOpenDocuments();
+            Assert.Contains(editedDocuments, static document => document.Path.EndsWith("open-saved.ps1") && !document.Saved);
+
+            MarkAsSaved(openSaved);
+            MarkAsSaved(openUntitled);
+
+            WorkspaceOpenDocument[] savedDocuments = editorOperationsService.GetWorkspaceOpenDocuments();
+            Assert.Contains(savedDocuments, static document => document.Path.EndsWith("open-saved.ps1") && document.Saved);
+            Assert.Contains(savedDocuments, static document => document.Path.StartsWith("untitled:", StringComparison.Ordinal) && !document.Saved);
+        }
+
+        private static ScriptFile CreateFileBuffer(WorkspaceService workspaceService, string fileName)
+        {
+            string filePath = Path.Combine(Path.GetTempPath(), Guid.NewGuid().ToString("N"), fileName);
+            return workspaceService.GetFileBuffer(DocumentUri.FromFileSystemPath(filePath), initialBuffer: string.Empty);
+        }
+
+        private static void MarkAsSaved(ScriptFile scriptFile) => scriptFile.IsInMemory = scriptFile.IsUntitled;
+    }
+}