Merge pull request #936 from aws-powertools/fix/override-lambda-console

chore: Console output overrides lambda console out

Author: hjgraca

libraries/src/AWS.Lambda.Powertools.Common/Core/ConsoleWrapper.cs

@@ -8,7 +8,6 @@ public class ConsoleWrapper : IConsoleWrapper
 {
     private static bool _override;
     private static TextWriter _testOutputStream;
-    private static bool _outputResetPerformed = false;
     private static bool _inTestMode = false;
 
     /// <inheritdoc />
@@ -20,7 +19,7 @@ public void WriteLine(string message)
         }
         else
         {
-            EnsureConsoleOutputOnce();
+            EnsureConsoleOutput();
             Console.WriteLine(message);
         }
     }
@@ -34,7 +33,7 @@ public void Debug(string message)
         }
         else
         {
-            EnsureConsoleOutputOnce();
+            EnsureConsoleOutput();
             System.Diagnostics.Debug.WriteLine(message);
         }
     }
@@ -50,9 +49,9 @@ public void Error(string message)
         {
             if (!_override)
             {
-                var errordOutput = new StreamWriter(Console.OpenStandardError());
-                errordOutput.AutoFlush = true;
-                Console.SetError(errordOutput);
+                var errorOutput = new StreamWriter(Console.OpenStandardError());
+                errorOutput.AutoFlush = true;
+                Console.SetError(errorOutput);
             }
             Console.Error.WriteLine(message);
         }
@@ -70,23 +69,69 @@ public static void SetOut(TextWriter consoleOut)
         Console.SetOut(consoleOut);
     }
 
-    private static void EnsureConsoleOutputOnce()
+    private static void EnsureConsoleOutput()
     {
-        if (_outputResetPerformed) return;
-        OverrideLambdaLogger();
-        _outputResetPerformed = true;
+        // Check if we need to override console output for Lambda environment
+        if (ShouldOverrideConsole())
+        {
+            OverrideLambdaLogger();
+        }
+    }
+
+    private static bool ShouldOverrideConsole()
+    {
+        // Don't override if we're in test mode
+        if (_inTestMode) return false;
+
+        // Always override in Lambda environment to prevent Lambda's log wrapping
+        var isLambda = !string.IsNullOrEmpty(Environment.GetEnvironmentVariable("AWS_LAMBDA_FUNCTION_NAME"));
+
+        return isLambda && (!_override || HasLambdaReInterceptedConsole());
+    }
+
+    internal static bool HasLambdaReInterceptedConsole()
+    {
+        return HasLambdaReInterceptedConsole(() => Console.Out);
     }
 
-    private static void OverrideLambdaLogger()
+    internal static bool HasLambdaReInterceptedConsole(Func<TextWriter> consoleOutAccessor)
     {
-        if (_override)
+        // Lambda might re-intercept console between init and handler execution
+        try
+        {
+            var currentOut = consoleOutAccessor();
+            // Check if current output stream looks like it might be Lambda's wrapper
+            var typeName = currentOut.GetType().FullName ?? "";
+            return typeName.Contains("Lambda") || typeName == "System.IO.TextWriter+SyncTextWriter";
+        }
+        catch
+        {
+            return true; // Assume re-interception if we can't determine
+        }
+    }
+    
+    internal static void OverrideLambdaLogger()
+    {
+        OverrideLambdaLogger(() => Console.OpenStandardOutput());
+    }
+    
+    internal static void OverrideLambdaLogger(Func<Stream> standardOutputOpener)
+    {
+        try
+        {
+            // Force override of LambdaLogger
+            var standardOutput = new StreamWriter(standardOutputOpener())
+            {
+                AutoFlush = true
+            };
+            Console.SetOut(standardOutput);
+            _override = true;
+        }
+        catch (Exception)
         {
-            return;
+            // Log the failure but don't throw - degraded functionality is better than crash
+            _override = false;
         }
-        // Force override of LambdaLogger
-        var standardOutput = new StreamWriter(Console.OpenStandardOutput());
-        standardOutput.AutoFlush = true;
-        Console.SetOut(standardOutput);
     }
 
     internal static void WriteLine(string logLevel, string message)
@@ -102,14 +147,14 @@ public static void ResetForTest()
         _override = false;
         _inTestMode = false;
         _testOutputStream = null;
-        _outputResetPerformed = false;
     }
 
     /// <summary>
     ///     Clear the output reset flag
     /// </summary>
     public static void ClearOutputResetFlag()
     {
-        _outputResetPerformed = false;
+        // This method is kept for backward compatibility but no longer needed
+        // since we removed the _outputResetPerformed flag
     }
 }