diff --git a/PowerShell.MCP.Proxy/PowerShell.MCP.Proxy.csproj b/PowerShell.MCP.Proxy/PowerShell.MCP.Proxy.csproj index 2484b38..804f937 100644 --- a/PowerShell.MCP.Proxy/PowerShell.MCP.Proxy.csproj +++ b/PowerShell.MCP.Proxy/PowerShell.MCP.Proxy.csproj @@ -61,4 +61,5 @@ + diff --git a/PowerShell.MCP/OutputTruncationHelper.cs b/PowerShell.MCP/OutputTruncationHelper.cs new file mode 100644 index 0000000..54ea749 --- /dev/null +++ b/PowerShell.MCP/OutputTruncationHelper.cs @@ -0,0 +1,173 @@ +using System.Diagnostics; + +namespace PowerShell.MCP; + +/// +/// Truncates large command output to preserve AI context window budget. +/// Saves the full output to a temp file so the AI can retrieve it via Read tool if needed. +/// +public static class OutputTruncationHelper +{ + internal const int TruncationThreshold = 15_000; + internal const int PreviewHeadSize = 1000; + internal const int PreviewTailSize = 1000; + + // The threshold must exceed the combined preview sizes; otherwise the head + // and tail slices overlap, producing duplicated content in the preview. + private static readonly bool _ = Validate(); + private static bool Validate() + { + Debug.Assert( + TruncationThreshold > PreviewHeadSize + PreviewTailSize, + $"TruncationThreshold ({TruncationThreshold}) must be greater than " + + $"PreviewHeadSize + PreviewTailSize ({PreviewHeadSize + PreviewTailSize}) " + + "to avoid overlapping head/tail previews."); + return true; + } + internal const string OutputDirectoryName = "PowerShell.MCP.Output"; + internal const int MaxFileAgeMinutes = 120; + internal const int NewlineScanLimit = 200; + + /// + /// Returns the output unchanged if within threshold, otherwise saves the full content + /// to a temp file and returns a head+tail preview with the file path. + /// + public static string TruncateIfNeeded(string output, string? outputDirectory = null) + { + if (output.Length <= TruncationThreshold) + return output; + + // Compute newline-aligned head boundary + var headEnd = FindHeadBoundary(output, PreviewHeadSize); + var head = output[..headEnd]; + + // Compute newline-aligned tail boundary + var tailStart = FindTailBoundary(output, PreviewTailSize); + var tail = output[tailStart..]; + + var omitted = output.Length - head.Length - tail.Length; + + var filePath = SaveOutputToFile(output, outputDirectory); + + var sb = new System.Text.StringBuilder(); + + if (filePath != null) + { + sb.AppendLine($"Output too large ({output.Length} characters). Full output saved to: {filePath}"); + sb.AppendLine($"Use invoke_expression('Show-TextFiles \"{filePath}\" -Contains \"search term\"') or -Pattern \"regex\" to search the output."); + } + else + { + // Disk save failed — still provide the preview without a file path + sb.AppendLine($"Output too large ({output.Length} characters). Could not save full output to file."); + } + + sb.AppendLine(); + sb.AppendLine("--- Preview (first ~1000 chars) ---"); + sb.AppendLine(head); + sb.AppendLine($"--- truncated ({omitted} chars omitted) ---"); + sb.AppendLine("--- Preview (last ~1000 chars) ---"); + sb.Append(tail); + + return sb.ToString(); + } + + /// + /// Saves the full output to a timestamped temp file and triggers opportunistic cleanup. + /// Returns the file path on success, null on failure. + /// + internal static string? SaveOutputToFile(string output, string? outputDirectory = null) + { + try + { + var directory = outputDirectory + ?? Path.Combine(Path.GetTempPath(), OutputDirectoryName); + + Directory.CreateDirectory(directory); + + var timestamp = DateTime.Now.ToString("yyyyMMdd_HHmmss_fff"); + var random = Path.GetRandomFileName(); + var fileName = $"pwsh_output_{timestamp}_{random}.txt"; + var filePath = Path.Combine(directory, fileName); + + File.WriteAllText(filePath, output); + + // Opportunistic cleanup — never let it block or fail the save + CleanupOldOutputFiles(directory); + + return filePath; + } + catch + { + return null; + } + } + + /// + /// Deletes pwsh_output_*.txt files older than minutes. + /// Each deletion is individually guarded so a locked file does not prevent other cleanups. + /// + internal static void CleanupOldOutputFiles(string? directory = null) + { + try + { + var dir = directory + ?? Path.Combine(Path.GetTempPath(), OutputDirectoryName); + + if (!Directory.Exists(dir)) + return; + + var cutoff = DateTime.Now.AddMinutes(-MaxFileAgeMinutes); + + foreach (var file in Directory.EnumerateFiles(dir, "pwsh_output_*.txt")) + { + try + { + if (File.GetLastWriteTime(file) < cutoff) + File.Delete(file); + } + catch (IOException) + { + // Another thread may be writing — safe to ignore + } + } + } + catch + { + // Directory enumeration itself failed — nothing to clean up + } + } + + /// + /// Finds a head cut position aligned to the nearest preceding newline within scan limit. + /// + private static int FindHeadBoundary(string output, int nominalSize) + { + if (nominalSize >= output.Length) + return output.Length; + + // Search backward from nominalSize for a newline, up to NewlineScanLimit chars + var searchStart = Math.Max(0, nominalSize - NewlineScanLimit); + var lastNewline = output.LastIndexOf('\n', nominalSize - 1, nominalSize - searchStart); + + // Cut after the newline to keep complete lines in the head + return lastNewline >= 0 ? lastNewline + 1 : nominalSize; + } + + /// + /// Finds a tail start position aligned to the nearest following newline within scan limit. + /// + private static int FindTailBoundary(string output, int nominalSize) + { + var nominalStart = output.Length - nominalSize; + if (nominalStart <= 0) + return 0; + + // Search forward from nominalStart for a newline, up to NewlineScanLimit chars + var searchEnd = Math.Min(output.Length, nominalStart + NewlineScanLimit); + var nextNewline = output.IndexOf('\n', nominalStart, searchEnd - nominalStart); + + // Start at the character after the newline to begin on a fresh line + return nextNewline >= 0 ? nextNewline + 1 : nominalStart; + } +} diff --git a/PowerShell.MCP/PowerShell.MCP.csproj b/PowerShell.MCP/PowerShell.MCP.csproj index 30e874e..e9ea058 100644 --- a/PowerShell.MCP/PowerShell.MCP.csproj +++ b/PowerShell.MCP/PowerShell.MCP.csproj @@ -8,6 +8,10 @@ $(NoWarn);NETSDK1206 + + + + diff --git a/PowerShell.MCP/Resources/MCPPollingEngine.ps1 b/PowerShell.MCP/Resources/MCPPollingEngine.ps1 index 9197dc2..66b3533 100644 --- a/PowerShell.MCP/Resources/MCPPollingEngine.ps1 +++ b/PowerShell.MCP/Resources/MCPPollingEngine.ps1 @@ -534,7 +534,7 @@ if (-not (Test-Path Variable:global:McpTimer)) { if ($null -eq $mcpOutput) { $mcpOutput = "Command execution completed" } - [PowerShell.MCP.Services.PowerShellCommunication]::NotifyResultReady($mcpOutput) + [PowerShell.MCP.Services.PowerShellCommunication]::NotifySilentResultReady($mcpOutput) } } } | Out-Null diff --git a/PowerShell.MCP/Services/PowerShellCommunication.cs b/PowerShell.MCP/Services/PowerShellCommunication.cs index a1e8113..4e7e3b8 100644 --- a/PowerShell.MCP/Services/PowerShellCommunication.cs +++ b/PowerShell.MCP/Services/PowerShellCommunication.cs @@ -17,7 +17,18 @@ public static void NotifyResultReady(string result) // Capture cache flag before AddToCache resets it _resultShouldCache = ExecutionState.ShouldCacheOutput; - // Always add to cache + // Truncate large output before caching to reduce pipe transfer and memory overhead + var output = OutputTruncationHelper.TruncateIfNeeded(result); + ExecutionState.AddToCache(output); + ExecutionState.CompleteExecution(); + _resultReadyEvent.Set(); + } + + /// + /// Silent command completed — no truncation (internal use, small known outputs) + /// + public static void NotifySilentResultReady(string result) + { ExecutionState.AddToCache(result); ExecutionState.CompleteExecution(); _resultReadyEvent.Set(); diff --git a/Tests/Unit/Core/OutputTruncationHelperTests.cs b/Tests/Unit/Core/OutputTruncationHelperTests.cs new file mode 100644 index 0000000..7a7173d --- /dev/null +++ b/Tests/Unit/Core/OutputTruncationHelperTests.cs @@ -0,0 +1,273 @@ +using PowerShell.MCP; +using Xunit; + +namespace PowerShell.MCP.Tests.Unit.Core; + +public class OutputTruncationHelperTests : IDisposable +{ + private readonly string _testDir; + + public OutputTruncationHelperTests() + { + _testDir = Path.Combine(Path.GetTempPath(), $"OutputTruncationHelperTests_{Guid.NewGuid():N}"); + Directory.CreateDirectory(_testDir); + } + + public void Dispose() + { + try { Directory.Delete(_testDir, recursive: true); } + catch { /* best-effort cleanup */ } + } + + #region TruncateIfNeeded Tests + + [Fact] + public void TruncateIfNeeded_Empty_ReturnsEmpty() + { + var result = OutputTruncationHelper.TruncateIfNeeded("", _testDir); + Assert.Equal("", result); + } + + [Fact] + public void TruncateIfNeeded_UnderThreshold_ReturnsUnchanged() + { + var output = new string('A', OutputTruncationHelper.TruncationThreshold - 1); + var result = OutputTruncationHelper.TruncateIfNeeded(output, _testDir); + Assert.Equal(output, result); + } + + [Fact] + public void TruncateIfNeeded_ExactThreshold_ReturnsUnchanged() + { + var output = new string('B', OutputTruncationHelper.TruncationThreshold); + var result = OutputTruncationHelper.TruncateIfNeeded(output, _testDir); + Assert.Equal(output, result); + } + + [Fact] + public void TruncateIfNeeded_OverThreshold_ReturnsTruncatedWithFilePath() + { + var output = new string('C', OutputTruncationHelper.TruncationThreshold + 1); + var result = OutputTruncationHelper.TruncateIfNeeded(output, _testDir); + + // Should contain file path info + Assert.Contains("Full output saved to:", result); + Assert.Contains("pwsh_output_", result); + Assert.Contains(".txt", result); + } + + [Fact] + public void TruncateIfNeeded_OverThreshold_SavesFullContentToFile() + { + var output = new string('D', OutputTruncationHelper.TruncationThreshold + 500); + OutputTruncationHelper.TruncateIfNeeded(output, _testDir); + + // Verify file was saved with full content + var files = Directory.GetFiles(_testDir, "pwsh_output_*.txt"); + Assert.Single(files); + var savedContent = File.ReadAllText(files[0]); + Assert.Equal(output, savedContent); + } + + [Fact] + public void TruncateIfNeeded_OverThreshold_PreviewContainsHeadAndTail() + { + // Build output with identifiable head and tail regions + var head = "HEAD_MARKER_" + new string('X', 500); + var middle = new string('M', OutputTruncationHelper.TruncationThreshold); + var tail = new string('Y', 500) + "_TAIL_MARKER"; + var output = head + middle + tail; + + var result = OutputTruncationHelper.TruncateIfNeeded(output, _testDir); + + Assert.Contains("HEAD_MARKER_", result); + Assert.Contains("_TAIL_MARKER", result); + Assert.Contains("truncated", result); + } + + [Fact] + public void TruncateIfNeeded_OverThreshold_PreviewTailMatchesOutputEnd() + { + // The tail preview should end with the same content as the original output + var tailContent = "FINAL_LINE_OF_OUTPUT"; + var output = new string('Z', OutputTruncationHelper.TruncationThreshold + 2000) + tailContent; + + var result = OutputTruncationHelper.TruncateIfNeeded(output, _testDir); + + Assert.EndsWith(tailContent, result); + } + + [Fact] + public void TruncateIfNeeded_OverThreshold_HeadAlignsToNewline() + { + // Place a newline within the scan range before PreviewHeadSize + // Head should cut at the newline boundary rather than at exactly PreviewHeadSize + var headSize = OutputTruncationHelper.PreviewHeadSize; + var newlinePos = headSize - 50; // newline within scan range + + var chars = new char[OutputTruncationHelper.TruncationThreshold + 2000]; + Array.Fill(chars, 'A'); + chars[newlinePos] = '\n'; + var output = new string(chars); + + var result = OutputTruncationHelper.TruncateIfNeeded(output, _testDir); + + // Extract the head content: appears after "--- Preview (first ~1000 chars) ---" + newline + var headMarker = "--- Preview (first ~1000 chars) ---" + Environment.NewLine; + var headStart = result.IndexOf(headMarker); + Assert.True(headStart >= 0, "Should contain head preview marker"); + var headContentStart = headStart + headMarker.Length; + + var truncatedMarker = "--- truncated"; + var truncatedIndex = result.IndexOf(truncatedMarker, headContentStart); + Assert.True(truncatedIndex > 0, "Should contain 'truncated' marker"); + + // Head content is between the marker and the truncated line (minus trailing newline from AppendLine) + var headContent = result.Substring(headContentStart, truncatedIndex - headContentStart) + .TrimEnd('\n', '\r'); + // Head should have been cut at the newline (position 950+1 = 951 chars) + Assert.Equal(newlinePos + 1, headContent.Length + 1); // +1 because the \n is the cut point + } + + [Fact] + public void TruncateIfNeeded_OverThreshold_TailAlignsToNewline() + { + // Place a newline near the tail start position + var tailSize = OutputTruncationHelper.PreviewTailSize; + var outputLength = OutputTruncationHelper.TruncationThreshold + 2000; + var tailStartPos = outputLength - tailSize; + var newlinePos = tailStartPos + 50; // newline shortly after tail start, within scan range + + var chars = new char[outputLength]; + Array.Fill(chars, 'B'); + chars[newlinePos] = '\n'; + var output = new string(chars); + + var result = OutputTruncationHelper.TruncateIfNeeded(output, _testDir); + + // Extract the tail content: appears after "--- Preview (last ~1000 chars) ---" + newline + var tailMarker = "--- Preview (last ~1000 chars) ---" + Environment.NewLine; + var tailMarkerPos = result.LastIndexOf(tailMarker); + Assert.True(tailMarkerPos >= 0, "Should contain tail preview marker"); + var tailContent = result.Substring(tailMarkerPos + tailMarker.Length); + + // Tail should start at the character after the planted newline + var expectedTail = output[(newlinePos + 1)..]; + Assert.Equal(expectedTail, tailContent); + } + + [Fact] + public void TruncateIfNeeded_OverThreshold_NoNewlineInScanRange_HardCuts() + { + // Output with no newlines at all — should hard-cut at exact positions + var output = new string('Q', OutputTruncationHelper.TruncationThreshold + 3000); + + var result = OutputTruncationHelper.TruncateIfNeeded(output, _testDir); + + // Should still produce a valid truncated result with head and tail + Assert.Contains("truncated", result); + Assert.Contains("Preview", result); + // The total result should be much shorter than the original output + Assert.True(result.Length < output.Length, "Truncated result should be shorter than original"); + } + + [Fact] + public void TruncateIfNeeded_FileSaveFails_StillReturnsPreview() + { + // Null characters are universally invalid in file paths across all platforms + var invalidDir = Path.Combine(_testDir, "invalid\0dir"); + + var output = new string('E', OutputTruncationHelper.TruncationThreshold + 1000); + var result = OutputTruncationHelper.TruncateIfNeeded(output, invalidDir); + + // Should still return a preview even though file save failed + Assert.Contains("truncated", result); + Assert.Contains("Preview", result); + // Should NOT contain "saved to" since the save failed + Assert.DoesNotContain("Full output saved to:", result); + } + + [Fact] + public void TruncateIfNeeded_CustomOutputDirectory_SavesToSpecifiedDir() + { + var customDir = Path.Combine(_testDir, "custom_output"); + Directory.CreateDirectory(customDir); + + var output = new string('F', OutputTruncationHelper.TruncationThreshold + 500); + var result = OutputTruncationHelper.TruncateIfNeeded(output, customDir); + + // File should be saved in the custom directory + var files = Directory.GetFiles(customDir, "pwsh_output_*.txt"); + Assert.Single(files); + Assert.Contains(customDir.Replace("\\", "/"), result.Replace("\\", "/")); + } + + [Fact] + public void TruncateIfNeeded_MessageShowsCharacters_NotKB() + { + var output = new string('G', OutputTruncationHelper.TruncationThreshold + 1000); + var result = OutputTruncationHelper.TruncateIfNeeded(output, _testDir); + + Assert.Contains("characters", result); + Assert.DoesNotContain("KB", result); + Assert.DoesNotContain("bytes", result); + } + + #endregion + + #region CleanupOldOutputFiles Tests + + [Fact] + public void CleanupOldOutputFiles_RemovesOldFiles() + { + // Create a file and set its last write time to beyond MaxFileAgeMinutes + var oldFile = Path.Combine(_testDir, "pwsh_output_old_test.txt"); + File.WriteAllText(oldFile, "old content"); + File.SetLastWriteTime(oldFile, DateTime.Now.AddMinutes(-(OutputTruncationHelper.MaxFileAgeMinutes + 10))); + + OutputTruncationHelper.CleanupOldOutputFiles(_testDir); + + Assert.False(File.Exists(oldFile), "Old file should have been deleted"); + } + + [Fact] + public void CleanupOldOutputFiles_KeepsRecentFiles() + { + // Create a recent file (within MaxFileAgeMinutes) + var recentFile = Path.Combine(_testDir, "pwsh_output_recent_test.txt"); + File.WriteAllText(recentFile, "recent content"); + // Last write time is now (default), well within the threshold + + OutputTruncationHelper.CleanupOldOutputFiles(_testDir); + + Assert.True(File.Exists(recentFile), "Recent file should NOT have been deleted"); + } + + [Fact] + public void CleanupOldOutputFiles_NonexistentDir_NoThrow() + { + var nonexistentDir = Path.Combine(_testDir, "does_not_exist_" + Guid.NewGuid().ToString("N")); + + // Should not throw even when directory doesn't exist + var exception = Record.Exception(() => OutputTruncationHelper.CleanupOldOutputFiles(nonexistentDir)); + Assert.Null(exception); + } + + [Fact] + public void CleanupOldOutputFiles_IgnoresIOException() + { + // Create an old file + var lockedFile = Path.Combine(_testDir, "pwsh_output_locked_test.txt"); + File.WriteAllText(lockedFile, "locked content"); + File.SetLastWriteTime(lockedFile, DateTime.Now.AddMinutes(-(OutputTruncationHelper.MaxFileAgeMinutes + 10))); + + // Lock the file by opening it with exclusive access + using var stream = new FileStream(lockedFile, FileMode.Open, FileAccess.Read, FileShare.None); + + // Should not throw even though the file is locked and cannot be deleted + var exception = Record.Exception(() => OutputTruncationHelper.CleanupOldOutputFiles(_testDir)); + Assert.Null(exception); + } + + #endregion +}