Skip to content

Commit 87022dc

Browse files
rjmholtrjmholt
committed
Fix remote psedit registration
1 parent 27a1735 commit 87022dc

File tree

1 file changed

+137
-103
lines changed

1 file changed

+137
-103
lines changed

‎src/PowerShellEditorServices/Services/Workspace/RemoteFileManagerService.cs‎

Lines changed: 137 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
using Microsoft.PowerShell.EditorServices.Services.PowerShell;
99
using Microsoft.PowerShell.EditorServices.Services.PowerShell.Execution;
1010
using Microsoft.PowerShell.EditorServices.Services.PowerShell.Runspace;
11+
using Microsoft.PowerShell.EditorServices.Services.PowerShell.Utility;
12+
using Microsoft.PowerShell.EditorServices.Utility;
1113
using System;
1214
using System.Collections.Generic;
1315
using System.Diagnostics;
@@ -18,6 +20,7 @@
1820
using System.Text;
1921
using System.Threading;
2022
using System.Threading.Tasks;
23+
using SMA = System.Management.Automation;
2124

2225
namespace Microsoft.PowerShell.EditorServices.Services
2326
{
@@ -258,7 +261,7 @@ public RemoteFileManagerService(
258261
this.logger = factory.CreateLogger<RemoteFileManagerService>();
259262
_runspaceContext = runspaceContext;
260263
_executionService = executionService;
261-
_executionService.RunspaceChanged += HandleRunspaceChangedAsync;
264+
_executionService.RunspaceChanged += HandleRunspaceChanged;
262265

263266
this.editorOperations = editorOperations;
264267

@@ -274,7 +277,7 @@ public RemoteFileManagerService(
274277

275278
// TODO: Do this somewhere other than the constructor and make it async
276279
// Register the psedit function in the current runspace
277-
this.RegisterPSEditFunction(_runspaceContext.CurrentRunspace);
280+
this.RegisterPSEditFunctionAsync().HandleErrorsAsync(logger);
278281
}
279282

280283
#endregion
@@ -505,10 +508,9 @@ private void StoreRemoteFile(
505508

506509
private RemotePathMappings GetPathMappings(IRunspaceInfo runspaceInfo)
507510
{
508-
RemotePathMappings remotePathMappings = null;
509511
string computerName = runspaceInfo.SessionDetails.ComputerName;
510512

511-
if (!this.filesPerComputer.TryGetValue(computerName, out remotePathMappings))
513+
if (!this.filesPerComputer.TryGetValue(computerName, out RemotePathMappings remotePathMappings))
512514
{
513515
remotePathMappings = new RemotePathMappings(runspaceInfo, this);
514516
this.filesPerComputer.Add(computerName, remotePathMappings);
@@ -517,28 +519,33 @@ private RemotePathMappings GetPathMappings(IRunspaceInfo runspaceInfo)
517519
return remotePathMappings;
518520
}
519521

520-
private async void HandleRunspaceChangedAsync(object sender, RunspaceChangedEventArgs e)
522+
private void HandleRunspaceChanged(object sender, RunspaceChangedEventArgs e)
521523
{
522524
if (e.ChangeAction == RunspaceChangeAction.Enter)
523525
{
524-
this.RegisterPSEditFunction(e.NewRunspace);
526+
this.RegisterPSEditFunction(e.NewRunspace.Runspace);
525527
return;
526528
}
527529

528530
// Close any remote files that were opened
529-
if (e.PreviousRunspace.IsOnRemoteMachine &&
530-
(e.ChangeAction == RunspaceChangeAction.Shutdown ||
531-
!string.Equals(
532-
e.NewRunspace.SessionDetails.ComputerName,
533-
e.PreviousRunspace.SessionDetails.ComputerName,
534-
StringComparison.CurrentCultureIgnoreCase)))
531+
if (ShouldTearDownRemoteFiles(e))
535532
{
536533
RemotePathMappings remotePathMappings;
537534
if (this.filesPerComputer.TryGetValue(e.PreviousRunspace.SessionDetails.ComputerName, out remotePathMappings))
538535
{
536+
var fileCloseTasks = new List<Task>();
539537
foreach (string remotePath in remotePathMappings.OpenedPaths)
540538
{
541-
await (this.editorOperations?.CloseFileAsync(remotePath)).ConfigureAwait(false);
539+
fileCloseTasks.Add(this.editorOperations?.CloseFileAsync(remotePath));
540+
}
541+
542+
try
543+
{
544+
Task.WaitAll(fileCloseTasks.ToArray());
545+
}
546+
catch (Exception ex)
547+
{
548+
this.logger.LogError(ex, "Unable to close all files in closed runspace");
542549
}
543550
}
544551
}
@@ -549,139 +556,166 @@ private async void HandleRunspaceChangedAsync(object sender, RunspaceChangedEven
549556
}
550557
}
551558

559+
private static bool ShouldTearDownRemoteFiles(RunspaceChangedEventArgs runspaceChangedEvent)
560+
{
561+
if (!runspaceChangedEvent.PreviousRunspace.IsOnRemoteMachine)
562+
{
563+
return false;
564+
}
565+
566+
if (runspaceChangedEvent.ChangeAction == RunspaceChangeAction.Shutdown)
567+
{
568+
return true;
569+
}
570+
571+
// Check to see if the runspace we're changing to is on a different machine to the one we left
572+
return !string.Equals(
573+
runspaceChangedEvent.NewRunspace.SessionDetails.ComputerName,
574+
runspaceChangedEvent.PreviousRunspace.SessionDetails.ComputerName,
575+
StringComparison.CurrentCultureIgnoreCase);
576+
}
577+
552578
private async void HandlePSEventReceivedAsync(object sender, PSEventArgs args)
553579
{
554-
if (string.Equals(RemoteSessionOpenFile, args.SourceIdentifier, StringComparison.CurrentCultureIgnoreCase))
580+
if (!string.Equals(RemoteSessionOpenFile, args.SourceIdentifier, StringComparison.CurrentCultureIgnoreCase))
555581
{
556-
try
582+
return;
583+
}
584+
585+
try
586+
{
587+
if (args.SourceArgs.Length >= 1)
557588
{
558-
if (args.SourceArgs.Length >= 1)
589+
string localFilePath = string.Empty;
590+
string remoteFilePath = args.SourceArgs[0] as string;
591+
592+
// Is this a local process runspace? Treat as a local file
593+
if (!_runspaceContext.CurrentRunspace.IsOnRemoteMachine)
559594
{
560-
string localFilePath = string.Empty;
561-
string remoteFilePath = args.SourceArgs[0] as string;
595+
localFilePath = remoteFilePath;
596+
}
597+
else
598+
{
599+
byte[] fileContent = null;
562600

563-
// Is this a local process runspace? Treat as a local file
564-
if (!_runspaceContext.CurrentRunspace.IsOnRemoteMachine)
565-
{
566-
localFilePath = remoteFilePath;
567-
}
568-
else
601+
if (args.SourceArgs.Length >= 2)
569602
{
570-
byte[] fileContent = null;
571-
572-
if (args.SourceArgs.Length >= 2)
603+
// Try to cast as a PSObject to get the BaseObject, if not, then try to case as a byte[]
604+
PSObject sourceObj = args.SourceArgs[1] as PSObject;
605+
if (sourceObj != null)
573606
{
574-
// Try to cast as a PSObject to get the BaseObject, if not, then try to case as a byte[]
575-
PSObject sourceObj = args.SourceArgs[1] as PSObject;
576-
if (sourceObj != null)
577-
{
578-
fileContent = sourceObj.BaseObject as byte[];
579-
}
580-
else
581-
{
582-
fileContent = args.SourceArgs[1] as byte[];
583-
}
584-
}
585-
586-
// If fileContent is still null after trying to
587-
// unpack the contents, just return an empty byte
588-
// array.
589-
fileContent = fileContent ?? Array.Empty<byte>();
590-
591-
if (remoteFilePath != null)
592-
{
593-
localFilePath =
594-
this.StoreRemoteFile(
595-
remoteFilePath,
596-
fileContent,
597-
_runspaceContext.CurrentRunspace);
607+
fileContent = sourceObj.BaseObject as byte[];
598608
}
599609
else
600610
{
601-
await (this.editorOperations?.NewFileAsync()).ConfigureAwait(false);
602-
EditorContext context = await (editorOperations?.GetEditorContextAsync()).ConfigureAwait(false);
603-
context?.CurrentFile.InsertText(Encoding.UTF8.GetString(fileContent, 0, fileContent.Length));
611+
fileContent = args.SourceArgs[1] as byte[];
604612
}
605613
}
606614

607-
bool preview = true;
608-
if (args.SourceArgs.Length >= 3)
615+
// If fileContent is still null after trying to
616+
// unpack the contents, just return an empty byte
617+
// array.
618+
fileContent = fileContent ?? Array.Empty<byte>();
619+
620+
if (remoteFilePath != null)
609621
{
610-
bool? previewCheck = args.SourceArgs[2] as bool?;
611-
preview = previewCheck ?? true;
622+
localFilePath =
623+
this.StoreRemoteFile(
624+
remoteFilePath,
625+
fileContent,
626+
_runspaceContext.CurrentRunspace);
612627
}
628+
else
629+
{
630+
await (this.editorOperations?.NewFileAsync()).ConfigureAwait(false);
631+
EditorContext context = await (editorOperations?.GetEditorContextAsync()).ConfigureAwait(false);
632+
context?.CurrentFile.InsertText(Encoding.UTF8.GetString(fileContent, 0, fileContent.Length));
633+
}
634+
}
613635

614-
// Open the file in the editor
615-
this.editorOperations?.OpenFileAsync(localFilePath, preview);
636+
bool preview = true;
637+
if (args.SourceArgs.Length >= 3)
638+
{
639+
bool? previewCheck = args.SourceArgs[2] as bool?;
640+
preview = previewCheck ?? true;
616641
}
617-
}
618-
catch (NullReferenceException e)
619-
{
620-
this.logger.LogException("Could not store null remote file content", e);
642+
643+
// Open the file in the editor
644+
await (this.editorOperations?.OpenFileAsync(localFilePath, preview)).ConfigureAwait(false);
621645
}
622646
}
647+
catch (NullReferenceException e)
648+
{
649+
this.logger.LogException("Could not store null remote file content", e);
650+
}
651+
catch (Exception e)
652+
{
653+
this.logger.LogException("Unable to handle remote file update", e);
654+
}
623655
}
624656

625-
private void RegisterPSEditFunction(IRunspaceInfo runspaceInfo)
657+
private Task RegisterPSEditFunctionAsync()
658+
=> _executionService.ExecuteDelegateAsync(
659+
"Register psedit function",
660+
ExecutionOptions.Default,
661+
(pwsh, cancellationToken) => RegisterPSEditFunction(pwsh.Runspace),
662+
CancellationToken.None);
663+
664+
private void RegisterPSEditFunction(Runspace runspace)
626665
{
627-
if (!runspaceInfo.IsOnRemoteMachine)
666+
if (!runspace.RunspaceIsRemote)
628667
{
629668
return;
630669
}
631670

632-
try
633-
{
634-
runspaceInfo.Runspace.Events.ReceivedEvents.PSEventReceived += HandlePSEventReceivedAsync;
671+
runspace.Events.ReceivedEvents.PSEventReceived += HandlePSEventReceivedAsync;
635672

636-
PSCommand createCommand = new PSCommand()
637-
.AddScript(CreatePSEditFunctionScript)
638-
.AddParameter("PSEditModule", PSEditModule);
673+
PSCommand createCommand = new PSCommand()
674+
.AddScript(CreatePSEditFunctionScript)
675+
.AddParameter("PSEditModule", PSEditModule);
639676

640-
if (runspaceInfo.RunspaceOrigin == RunspaceOrigin.DebuggedRunspace)
641-
{
642-
_executionService.ExecutePSCommandAsync(createCommand, CancellationToken.None).GetAwaiter().GetResult();
643-
}
644-
else
645-
{
646-
using (var powerShell = System.Management.Automation.PowerShell.Create())
647-
{
648-
powerShell.Runspace = runspaceInfo.Runspace;
649-
powerShell.Commands = createCommand;
650-
powerShell.Invoke();
651-
}
652-
}
677+
var pwsh = SMA.PowerShell.Create();
678+
pwsh.Runspace = runspace;
679+
try
680+
{
681+
pwsh.InvokeCommand(createCommand, new PSInvocationSettings { AddToHistory = false, ErrorActionPreference = ActionPreference.Stop });
653682
}
654-
catch (RemoteException e)
683+
catch (Exception e)
655684
{
656685
this.logger.LogException("Could not create psedit function.", e);
657686
}
687+
finally
688+
{
689+
pwsh.Dispose();
690+
}
658691
}
659692

660693
private void RemovePSEditFunction(IRunspaceInfo runspaceInfo)
661694
{
662-
if (runspaceInfo.RunspaceOrigin == RunspaceOrigin.PSSession)
695+
if (runspaceInfo.RunspaceOrigin != RunspaceOrigin.PSSession)
663696
{
664-
try
697+
return;
698+
}
699+
try
700+
{
701+
if (runspaceInfo.Runspace.Events != null)
665702
{
666-
if (runspaceInfo.Runspace.Events != null)
667-
{
668-
runspaceInfo.Runspace.Events.ReceivedEvents.PSEventReceived -= HandlePSEventReceivedAsync;
669-
}
703+
runspaceInfo.Runspace.Events.ReceivedEvents.PSEventReceived -= HandlePSEventReceivedAsync;
704+
}
670705

671-
if (runspaceInfo.Runspace.RunspaceStateInfo.State == RunspaceState.Opened)
706+
if (runspaceInfo.Runspace.RunspaceStateInfo.State == RunspaceState.Opened)
707+
{
708+
using (var powerShell = SMA.PowerShell.Create())
672709
{
673-
using (var powerShell = System.Management.Automation.PowerShell.Create())
674-
{
675-
powerShell.Runspace = runspaceInfo.Runspace;
676-
powerShell.Commands.AddScript(RemovePSEditFunctionScript);
677-
powerShell.Invoke();
678-
}
710+
powerShell.Runspace = runspaceInfo.Runspace;
711+
powerShell.Commands.AddScript(RemovePSEditFunctionScript);
712+
powerShell.Invoke();
679713
}
680714
}
681-
catch (Exception e) when (e is RemoteException || e is PSInvalidOperationException)
682-
{
683-
this.logger.LogException("Could not remove psedit function.", e);
684-
}
715+
}
716+
catch (Exception e) when (e is RemoteException || e is PSInvalidOperationException)
717+
{
718+
this.logger.LogException("Could not remove psedit function.", e);
685719
}
686720
}
687721

0 commit comments

Comments
 (0)