From 2f3acb19f4328d21fd46fe2d691583633dc5febb Mon Sep 17 00:00:00 2001
From: Beena352 <beenachauhan@microsoft.com>
Date: Mon, 22 Dec 2025 17:37:50 -0800
Subject: [PATCH 01/11] Refactor wsladiag CLI parsing and fix shell PTY
 handling

---
 src/shared/inc/CommandLine.h      |  12 +++-
 src/windows/wsladiag/wsladiag.cpp | 112 +++++++++++++++++++++---------
 2 files changed, 88 insertions(+), 36 deletions(-)

diff --git a/src/shared/inc/CommandLine.h b/src/shared/inc/CommandLine.h
index 9966c5e65..d53991846 100644
--- a/src/shared/inc/CommandLine.h
+++ b/src/shared/inc/CommandLine.h
@@ -307,7 +307,8 @@ class ArgumentParser
 public:
 #ifdef WIN32
 
-    ArgumentParser(const std::wstring& CommandLine, LPCWSTR Name, int StartIndex = 1) : m_startIndex(StartIndex), m_name(Name)
+    ArgumentParser(const std::wstring& CommandLine, LPCWSTR Name, int StartIndex = 1, bool IgnoreUnknownArgs = false) :
+        m_startIndex(StartIndex), m_name(Name), m_ignoreUnknownArgs(IgnoreUnknownArgs)
     {
         m_argv.reset(CommandLineToArgvW(std::wstring(CommandLine).c_str(), &m_argc));
         THROW_LAST_ERROR_IF(!m_argv);
@@ -315,7 +316,8 @@ class ArgumentParser
 
 #else
 
-    ArgumentParser(int argc, const char* const* argv) : m_argc(argc), m_argv(argv), m_startIndex(1)
+    ArgumentParser(int argc, const char* const* argv, bool IgnoreUnknownArgs = false) :
+        m_argc(argc), m_argv(argv), m_startIndex(1), m_ignoreUnknownArgs(IgnoreUnknownArgs)
     {
     }
 
@@ -429,6 +431,11 @@ class ArgumentParser
 
             if (!foundMatch)
             {
+                if (m_ignoreUnknownArgs)
+                {
+                    break;
+                }
+
                 THROW_USER_ERROR(wsl::shared::Localization::MessageInvalidCommandLine(m_argv[i], m_name ? m_name : m_argv[0]));
             }
 
@@ -535,6 +542,7 @@ class ArgumentParser
 
     int m_startIndex{};
     const TChar* m_name{};
+    bool m_ignoreUnknownArgs{false};
 };
 } // namespace wsl::shared
 
diff --git a/src/windows/wsladiag/wsladiag.cpp b/src/windows/wsladiag/wsladiag.cpp
index eef2f3150..23a00cce5 100644
--- a/src/windows/wsladiag/wsladiag.cpp
+++ b/src/windows/wsladiag/wsladiag.cpp
@@ -8,7 +8,7 @@ Module Name:
 
 Abstract:
 
-    Entry point for the wsladiag tool, performs WSL runtime initialization and parses --list/--help.
+    Entry point for the wsladiag tool, performs WSL runtime initialization and parses list/shell/help.
 
 --*/
 
@@ -42,9 +42,35 @@ static int ReportError(const std::wstring& context, HRESULT hr)
     return 1;
 }
 
-// Handler for `wsladiag shell <SessionName>` (TTY-backed interactive shell).
-static int RunShellCommand(const std::wstring& sessionName, bool verbose)
+// Handler for `wsladiag shell <SessionName> [--verbose]` command - launches TTY-backed interactive shell.
+static int RunShellCommand(std::wstring_view commandLine)
 {
+    std::wstring sessionName;
+    bool verbose = false;
+
+    ArgumentParser parser(std::wstring{commandLine}, L"wsladiag", 2); // Skip "wsladiag.exe shell" to parse shell-specific args
+    parser.AddPositionalArgument(sessionName, 0);
+    parser.AddArgument(verbose, L"--verbose", L'v');
+
+    try
+    {
+        parser.Parse();
+    }
+    catch (...)
+    {
+        const auto hr = wil::ResultFromCaughtException();
+        if (hr == E_INVALIDARG)
+        {
+            return 1;
+        }
+        throw;
+    }
+
+    if (sessionName.empty())
+    {
+        return 1;
+    }
+
     const auto log = [&](std::wstring_view msg) {
         if (verbose)
         {
@@ -85,9 +111,11 @@ static int RunShellCommand(const std::wstring& sessionName, bool verbose)
     wsl::windows::common::WSLAProcessLauncher launcher{
         shell, {shell, "--login"}, {"TERM=xterm-256color"}, wsl::windows::common::ProcessFlags::None};
 
-    launcher.AddFd(WSLA_PROCESS_FD{.Fd = 0, .Type = WSLAFdTypeTerminalInput});
-    launcher.AddFd(WSLA_PROCESS_FD{.Fd = 1, .Type = WSLAFdTypeTerminalOutput});
-    launcher.AddFd(WSLA_PROCESS_FD{.Fd = 2, .Type = WSLAFdTypeTerminalControl});
+    launcher.AddFd(WSLA_PROCESS_FD{.Fd = 0, .Type = WSLAFdTypeTerminalInput, .Path = nullptr});
+    launcher.AddFd(WSLA_PROCESS_FD{.Fd = 1, .Type = WSLAFdTypeTerminalOutput, .Path = nullptr});
+    launcher.AddFd(WSLA_PROCESS_FD{.Fd = 2, .Type = WSLAFdTypeTerminalControl, .Path = nullptr});
+
+    log(std::format(L"[diag] tty rows={} cols={}", rows, cols));
     launcher.SetTtySize(rows, cols);
 
     log(L"[diag] launching shell process...");
@@ -96,7 +124,6 @@ static int RunShellCommand(const std::wstring& sessionName, bool verbose)
 
     auto ttyIn = process.GetStdHandle(0);
     auto ttyOut = process.GetStdHandle(1);
-    auto ttyControl = process.GetStdHandle(2);
 
     // Console handles.
     wil::unique_hfile conin{
@@ -119,11 +146,11 @@ static int RunShellCommand(const std::wstring& sessionName, bool verbose)
     const UINT originalOutCP = GetConsoleOutputCP();
     const UINT originalInCP = GetConsoleCP();
 
-    auto restoreConsole = wil::scope_exit([&] {
-        SetConsoleMode(consoleIn, originalInMode);
-        SetConsoleMode(consoleOut, originalOutMode);
-        SetConsoleOutputCP(originalOutCP);
-        SetConsoleCP(originalInCP);
+    auto restoreConsole = wil::scope_exit_log(WI_DIAGNOSTICS_INFO, [&] {
+        LOG_IF_WIN32_BOOL_FALSE(SetConsoleMode(consoleIn, originalInMode));
+        LOG_IF_WIN32_BOOL_FALSE(SetConsoleMode(consoleOut, originalOutMode));
+        LOG_IF_WIN32_BOOL_FALSE(SetConsoleOutputCP(originalOutCP));
+        LOG_IF_WIN32_BOOL_FALSE(SetConsoleCP(originalInCP));
     });
 
     // Console mode for interactive terminal.
@@ -139,23 +166,46 @@ static int RunShellCommand(const std::wstring& sessionName, bool verbose)
     THROW_LAST_ERROR_IF(!SetConsoleOutputCP(CP_UTF8));
     THROW_LAST_ERROR_IF(!SetConsoleCP(CP_UTF8));
 
-    // Keep terminal control socket alive.
     auto exitEvent = wil::unique_event(wil::EventOptions::ManualReset);
-    wsl::shared::SocketChannel controlChannel{wil::unique_socket{(SOCKET)ttyControl.release()}, "TerminalControl", exitEvent.get()};
+
+    std::optional<wsl::shared::SocketChannel> controlChannel;
+    try
+    {
+        auto ttyControl = process.GetStdHandle(2); // only once
+        controlChannel.emplace(wil::unique_socket{(SOCKET)ttyControl.release()}, "TerminalControl", exitEvent.get());
+    }
+    catch (const wil::ResultException& e)
+    {
+        const HRESULT hr = e.GetErrorCode();
+        if (hr == HRESULT_FROM_WIN32(ERROR_NOT_FOUND))
+        {
+            log(L"[diag] TerminalControl not available; live resize disabled");
+        }
+        else
+        {
+            throw; // or ReportError + return
+        }
+    }
 
     auto updateTerminalSize = [&]() {
+        if (!controlChannel.has_value())
+        {
+            return;
+        }
+
         CONSOLE_SCREEN_BUFFER_INFOEX infoEx{};
         infoEx.cbSize = sizeof(infoEx);
         THROW_IF_WIN32_BOOL_FALSE(GetConsoleScreenBufferInfoEx(consoleOut, &infoEx));
 
         WSLA_TERMINAL_CHANGED message{};
-        message.Columns = infoEx.srWindow.Right - infoEx.srWindow.Left + 1;
-        message.Rows = infoEx.srWindow.Bottom - infoEx.srWindow.Top + 1;
+        message.Columns = static_cast<unsigned short>(infoEx.srWindow.Right - infoEx.srWindow.Left + 1);
+        message.Rows = static_cast<unsigned short>(infoEx.srWindow.Bottom - infoEx.srWindow.Top + 1);
 
-        controlChannel.SendMessage(message);
+        controlChannel->SendMessage(message);
     };
 
-    // Relay console -> tty input.
+    // Start input relay thread to forward console input to TTY
+    // Runs in parallel with output relay (main thread)
     std::thread inputThread([&] {
         try
         {
@@ -178,16 +228,19 @@ static int RunShellCommand(const std::wstring& sessionName, bool verbose)
     // Relay tty output -> console (blocks until output ends).
     wsl::windows::common::relay::InterruptableRelay(ttyOut.get(), consoleOut, exitEvent.get());
 
+    // Output relay finished - signal input thread to stop and wait for process exit
+    exitEvent.SetEvent();
     process.GetExitEvent().wait();
-    auto [code, signalled] = process.GetExitState();
+
+    auto [exitCode, signalled] = process.GetExitState();
 
     std::wstring shellWide(shell.begin(), shell.end());
-    wslutil::PrintMessage(std::format(L"{} exited with: {}{}", shellWide, code, signalled ? L" (signalled)" : L""), stdout);
+    wslutil::PrintMessage(std::format(L"{} exited with: {}{}", shellWide, exitCode, signalled ? L" (signalled)" : L""), stdout);
 
     return 0;
 }
 
-static int RunListCommand(bool /*verbose*/)
+static int RunListCommand()
 {
     wil::com_ptr<IWSLAUserSession> userSession;
     THROW_IF_FAILED(CoCreateInstance(__uuidof(WSLAUserSession), nullptr, CLSCTX_LOCAL_SERVER, IID_PPV_ARGS(&userSession)));
@@ -257,17 +310,13 @@ int wsladiag_main(std::wstring_view commandLine)
     THROW_IF_WIN32_ERROR(WSAStartup(MAKEWORD(2, 2), &data));
     auto wsaCleanup = wil::scope_exit_log(WI_DIAGNOSTICS_INFO, []() { WSACleanup(); });
 
-    ArgumentParser parser(std::wstring{commandLine}, L"wsladiag");
+    ArgumentParser parser(std::wstring{commandLine}, L"wsladiag", 1, true);
 
     bool help = false;
-    bool verbose = false;
     std::wstring verb;
-    std::wstring shellSession;
 
-    parser.AddPositionalArgument(verb, 0);         // "list" or "shell"
-    parser.AddPositionalArgument(shellSession, 1); // session name for "shell"
+    parser.AddPositionalArgument(verb, 0);
     parser.AddArgument(help, L"--help", L'h');
-    parser.AddArgument(verbose, L"--verbose", L'v');
 
     auto printUsage = []() {
         wslutil::PrintMessage(
@@ -301,16 +350,11 @@ int wsladiag_main(std::wstring_view commandLine)
     }
     else if (verb == L"list")
     {
-        return RunListCommand(verbose);
+        return RunListCommand();
     }
     else if (verb == L"shell")
     {
-        if (shellSession.empty())
-        {
-            printUsage();
-            return 1;
-        }
-        return RunShellCommand(shellSession, verbose);
+        return RunShellCommand(commandLine);
     }
     else
     {

From 68b43d2be7888587588580385d111c171ad35ca6 Mon Sep 17 00:00:00 2001
From: Beena352 <beenachauhan@microsoft.com>
Date: Tue, 23 Dec 2025 10:22:20 -0800
Subject: [PATCH 02/11] Rename ArgumentParser flag to stopOnUnknownArgs

---
 src/shared/inc/CommandLine.h      |  8 +--
 src/windows/wsladiag/wsladiag.cpp | 90 +++++++++----------------------
 2 files changed, 28 insertions(+), 70 deletions(-)

diff --git a/src/shared/inc/CommandLine.h b/src/shared/inc/CommandLine.h
index d53991846..bc242ecd5 100644
--- a/src/shared/inc/CommandLine.h
+++ b/src/shared/inc/CommandLine.h
@@ -308,7 +308,7 @@ class ArgumentParser
 #ifdef WIN32
 
     ArgumentParser(const std::wstring& CommandLine, LPCWSTR Name, int StartIndex = 1, bool IgnoreUnknownArgs = false) :
-        m_startIndex(StartIndex), m_name(Name), m_ignoreUnknownArgs(IgnoreUnknownArgs)
+        m_startIndex(StartIndex), m_name(Name), m_stopUnknownArgs(IgnoreUnknownArgs)
     {
         m_argv.reset(CommandLineToArgvW(std::wstring(CommandLine).c_str(), &m_argc));
         THROW_LAST_ERROR_IF(!m_argv);
@@ -317,7 +317,7 @@ class ArgumentParser
 #else
 
     ArgumentParser(int argc, const char* const* argv, bool IgnoreUnknownArgs = false) :
-        m_argc(argc), m_argv(argv), m_startIndex(1), m_ignoreUnknownArgs(IgnoreUnknownArgs)
+        m_argc(argc), m_argv(argv), m_startIndex(1), m_stopUnknownArgs(IgnoreUnknownArgs)
     {
     }
 
@@ -431,7 +431,7 @@ class ArgumentParser
 
             if (!foundMatch)
             {
-                if (m_ignoreUnknownArgs)
+                if (m_stopUnknownArgs)
                 {
                     break;
                 }
@@ -542,7 +542,7 @@ class ArgumentParser
 
     int m_startIndex{};
     const TChar* m_name{};
-    bool m_ignoreUnknownArgs{false};
+    bool m_stopUnknownArgs{false};
 };
 } // namespace wsl::shared
 
diff --git a/src/windows/wsladiag/wsladiag.cpp b/src/windows/wsladiag/wsladiag.cpp
index 23a00cce5..d785d6435 100644
--- a/src/windows/wsladiag/wsladiag.cpp
+++ b/src/windows/wsladiag/wsladiag.cpp
@@ -52,23 +52,11 @@ static int RunShellCommand(std::wstring_view commandLine)
     parser.AddPositionalArgument(sessionName, 0);
     parser.AddArgument(verbose, L"--verbose", L'v');
 
-    try
-    {
-        parser.Parse();
-    }
-    catch (...)
-    {
-        const auto hr = wil::ResultFromCaughtException();
-        if (hr == E_INVALIDARG)
-        {
-            return 1;
-        }
-        throw;
-    }
+    parser.Parse();
 
     if (sessionName.empty())
     {
-        return 1;
+        THROW_HR(E_INVALIDARG);
     }
 
     const auto log = [&](std::wstring_view msg) {
@@ -168,31 +156,10 @@ static int RunShellCommand(std::wstring_view commandLine)
 
     auto exitEvent = wil::unique_event(wil::EventOptions::ManualReset);
 
-    std::optional<wsl::shared::SocketChannel> controlChannel;
-    try
-    {
-        auto ttyControl = process.GetStdHandle(2); // only once
-        controlChannel.emplace(wil::unique_socket{(SOCKET)ttyControl.release()}, "TerminalControl", exitEvent.get());
-    }
-    catch (const wil::ResultException& e)
-    {
-        const HRESULT hr = e.GetErrorCode();
-        if (hr == HRESULT_FROM_WIN32(ERROR_NOT_FOUND))
-        {
-            log(L"[diag] TerminalControl not available; live resize disabled");
-        }
-        else
-        {
-            throw; // or ReportError + return
-        }
-    }
+    auto ttyControl = process.GetStdHandle(2); // TerminalControl
+    wsl::shared::SocketChannel controlChannel{wil::unique_socket{(SOCKET)ttyControl.release()}, "TerminalControl", exitEvent.get()};
 
     auto updateTerminalSize = [&]() {
-        if (!controlChannel.has_value())
-        {
-            return;
-        }
-
         CONSOLE_SCREEN_BUFFER_INFOEX infoEx{};
         infoEx.cbSize = sizeof(infoEx);
         THROW_IF_WIN32_BOOL_FALSE(GetConsoleScreenBufferInfoEx(consoleOut, &infoEx));
@@ -201,7 +168,7 @@ static int RunShellCommand(std::wstring_view commandLine)
         message.Columns = static_cast<unsigned short>(infoEx.srWindow.Right - infoEx.srWindow.Left + 1);
         message.Rows = static_cast<unsigned short>(infoEx.srWindow.Bottom - infoEx.srWindow.Top + 1);
 
-        controlChannel->SendMessage(message);
+        controlChannel.SendMessage(message);
     };
 
     // Start input relay thread to forward console input to TTY
@@ -228,8 +195,6 @@ static int RunShellCommand(std::wstring_view commandLine)
     // Relay tty output -> console (blocks until output ends).
     wsl::windows::common::relay::InterruptableRelay(ttyOut.get(), consoleOut, exitEvent.get());
 
-    // Output relay finished - signal input thread to stop and wait for process exit
-    exitEvent.SetEvent();
     process.GetExitEvent().wait();
 
     auto [exitCode, signalled] = process.GetExitState();
@@ -293,6 +258,17 @@ static int RunListCommand()
     return 0;
 }
 
+static void PrintUsage()
+{
+    wslutil::PrintMessage(
+        L"wsladiag - WSLA diagnostics tool\n"
+        L"Usage:\n"
+        L"  wsladiag list\n"
+        L"  wsladiag shell <SessionName> [--verbose]\n"
+        L"  wsladiag --help\n",
+        stderr);
+}
+
 int wsladiag_main(std::wstring_view commandLine)
 {
     wslutil::ConfigureCrt();
@@ -318,34 +294,11 @@ int wsladiag_main(std::wstring_view commandLine)
     parser.AddPositionalArgument(verb, 0);
     parser.AddArgument(help, L"--help", L'h');
 
-    auto printUsage = []() {
-        wslutil::PrintMessage(
-            L"wsladiag - WSLA diagnostics tool\n"
-            L"Usage:\n"
-            L"  wsladiag list\n"
-            L"  wsladiag shell <SessionName> [--verbose]\n"
-            L"  wsladiag --help\n",
-            stderr);
-    };
-
-    try
-    {
-        parser.Parse();
-    }
-    catch (...)
-    {
-        const auto hr = wil::ResultFromCaughtException();
-        if (hr == E_INVALIDARG)
-        {
-            printUsage();
-            return 1;
-        }
-        throw;
-    }
+    parser.Parse(); // Let exceptions propagate to wmain for centralized handling
 
     if (help || verb.empty())
     {
-        printUsage();
+        PrintUsage();
         return 0;
     }
     else if (verb == L"list")
@@ -359,7 +312,7 @@ int wsladiag_main(std::wstring_view commandLine)
     else
     {
         wslutil::PrintMessage(std::format(L"Unknown command: '{}'", verb), stderr);
-        printUsage();
+        PrintUsage();
         return 1;
     }
 }
@@ -373,6 +326,11 @@ int wmain(int, wchar_t**)
     catch (...)
     {
         const auto hr = wil::ResultFromCaughtException();
+        if (hr == E_INVALIDARG)
+        {
+            PrintUsage();
+            return 1;
+        }
         return ReportError(L"wsladiag failed", hr);
     }
 }

From 33542d760eeb9158c6526a0f43b33ed8019ee379 Mon Sep 17 00:00:00 2001
From: Beena352 <beenachauhan@microsoft.com>
Date: Wed, 24 Dec 2025 10:35:20 -0800
Subject: [PATCH 03/11] Rename ArgumentParser constructor flag to
 stopOnUnknownArgs

---
 src/shared/inc/CommandLine.h      | 10 +++++-----
 src/windows/wsladiag/wsladiag.cpp |  4 ++--
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/shared/inc/CommandLine.h b/src/shared/inc/CommandLine.h
index bc242ecd5..3ef367ed4 100644
--- a/src/shared/inc/CommandLine.h
+++ b/src/shared/inc/CommandLine.h
@@ -307,8 +307,8 @@ class ArgumentParser
 public:
 #ifdef WIN32
 
-    ArgumentParser(const std::wstring& CommandLine, LPCWSTR Name, int StartIndex = 1, bool IgnoreUnknownArgs = false) :
-        m_startIndex(StartIndex), m_name(Name), m_stopUnknownArgs(IgnoreUnknownArgs)
+    ArgumentParser(const std::wstring& CommandLine, LPCWSTR Name, int StartIndex = 1, bool stopUnknownArgs = false) :
+        m_startIndex(StartIndex), m_name(Name), m_stopUnknownArgs(stopUnknownArgs)
     {
         m_argv.reset(CommandLineToArgvW(std::wstring(CommandLine).c_str(), &m_argc));
         THROW_LAST_ERROR_IF(!m_argv);
@@ -316,8 +316,8 @@ class ArgumentParser
 
 #else
 
-    ArgumentParser(int argc, const char* const* argv, bool IgnoreUnknownArgs = false) :
-        m_argc(argc), m_argv(argv), m_startIndex(1), m_stopUnknownArgs(IgnoreUnknownArgs)
+    ArgumentParser(int argc, const char* const* argv, bool stopUnknownArgs = false) :
+        m_argc(argc), m_argv(argv), m_startIndex(1), m_stopUnknownArgs(stopUnknownArgs)
     {
     }
 
@@ -402,7 +402,7 @@ class ArgumentParser
                         const TChar* value = nullptr;
                         if (e.Positional)
                         {
-                            value = m_argv[i]; // Positional arguments directly receive arvg[i]
+                            value = m_argv[i]; // Positional arguments directly receive argv[i]
                         }
                         else if (i + 1 < m_argc)
                         {
diff --git a/src/windows/wsladiag/wsladiag.cpp b/src/windows/wsladiag/wsladiag.cpp
index d785d6435..36fa5db7e 100644
--- a/src/windows/wsladiag/wsladiag.cpp
+++ b/src/windows/wsladiag/wsladiag.cpp
@@ -165,8 +165,8 @@ static int RunShellCommand(std::wstring_view commandLine)
         THROW_IF_WIN32_BOOL_FALSE(GetConsoleScreenBufferInfoEx(consoleOut, &infoEx));
 
         WSLA_TERMINAL_CHANGED message{};
-        message.Columns = static_cast<unsigned short>(infoEx.srWindow.Right - infoEx.srWindow.Left + 1);
-        message.Rows = static_cast<unsigned short>(infoEx.srWindow.Bottom - infoEx.srWindow.Top + 1);
+        message.Columns = infoEx.srWindow.Right - infoEx.srWindow.Left + 1;
+        message.Rows = infoEx.srWindow.Bottom - infoEx.srWindow.Top + 1;
 
         controlChannel.SendMessage(message);
     };

From 72eea12c5bf9c77718cd14a13d1dbd81f48178ff Mon Sep 17 00:00:00 2001
From: Beena352 <beenachauhan@microsoft.com>
Date: Tue, 23 Dec 2025 12:59:51 -0800
Subject: [PATCH 04/11] WIP: local test changes

---
 test/windows/CMakeLists.txt | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/windows/CMakeLists.txt b/test/windows/CMakeLists.txt
index c700d66f4..1abe54e68 100644
--- a/test/windows/CMakeLists.txt
+++ b/test/windows/CMakeLists.txt
@@ -9,7 +9,8 @@ set(SOURCES
     PluginTests.cpp
     PolicyTests.cpp
     InstallerTests.cpp
-    WSLATests.cpp)
+    WSLATests.cpp
+    WsladiagTests.cpp)
 
 set(HEADERS
     Common.h

From 12877e8de66491bd28b4e64e88cad0e158a36ccb Mon Sep 17 00:00:00 2001
From: Beena352 <beenachauhan@microsoft.com>
Date: Tue, 23 Dec 2025 14:55:52 -0800
Subject: [PATCH 05/11] Add argument and error-handling tests for wsladiag

---
 test/windows/WsladiagTests.cpp | 134 +++++++++++++++++++++++++++++++++
 1 file changed, 134 insertions(+)
 create mode 100644 test/windows/WsladiagTests.cpp

diff --git a/test/windows/WsladiagTests.cpp b/test/windows/WsladiagTests.cpp
new file mode 100644
index 000000000..bd3f6801f
--- /dev/null
+++ b/test/windows/WsladiagTests.cpp
@@ -0,0 +1,134 @@
+/*++
+
+Copyright (c) Microsoft. All rights reserved.
+
+Module Name:
+
+    WsladiagTests.cpp
+
+Abstract:
+
+    This file contains smoke tests for wsladiag.
+
+--*/
+
+#include "precomp.h"
+#include "Common.h"
+#include <format>
+
+namespace WsladiagTests {
+class WsladiagTests
+{
+    WSL_TEST_CLASS(WsladiagTests)
+
+    TEST_CLASS_SETUP(TestClassSetup)
+    {
+        VERIFY_ARE_EQUAL(LxsstuInitialize(FALSE), TRUE);
+        return true;
+    }
+
+    TEST_CLASS_CLEANUP(TestClassCleanup)
+    {
+        LxsstuUninitialize(FALSE);
+        return true;
+    }
+
+    static std::wstring BuildWsladiagCmd(const std::wstring& args)
+    {
+        const auto exePath = wsl::windows::common::wslutil::GetBasePath() / L"wsladiag.exe";
+        VERIFY_IS_TRUE(std::filesystem::exists(exePath));
+
+        const auto exe = exePath.wstring();
+        return args.empty() ? std::format(L"\"{}\"", exe) : std::format(L"\"{}\" {}", exe, args);
+    }
+
+    static std::tuple<std::wstring, std::wstring, int> RunWsladiag(const std::wstring& args)
+    {
+        auto cmd = BuildWsladiagCmd(args);
+        return LxsstuLaunchCommandAndCaptureOutputWithResult(cmd.data());
+    }
+
+    TEST_METHOD(List_ShowsSessionsOrNoSessions)
+    {
+        auto [out, err, code] = RunWsladiag(L"list");
+        VERIFY_ARE_EQUAL(0, code);
+        VERIFY_ARE_EQUAL(L"", err);
+
+        const bool noSessions = (out.find(L"No WSLA sessions found.") != std::wstring::npos);
+
+        const bool hasTable = (out.find(L"WSLA session") != std::wstring::npos) && (out.find(L"ID") != std::wstring::npos) &&
+                              (out.find(L"Display Name") != std::wstring::npos);
+
+        VERIFY_IS_TRUE(noSessions || hasTable);
+    }
+
+    TEST_METHOD(Help_ShowsUsage)
+    {
+        auto [out, err, code] = RunWsladiag(L"--help");
+        VERIFY_ARE_EQUAL(0, code);
+        VERIFY_ARE_EQUAL(L"", out);
+
+        VERIFY_IS_TRUE(err.find(L"Usage:") != std::wstring::npos);
+        VERIFY_IS_TRUE(err.find(L"wsladiag list") != std::wstring::npos);
+        VERIFY_IS_TRUE(err.find(L"wsladiag shell <SessionName> [--verbose]") != std::wstring::npos);
+    }
+
+    TEST_METHOD(Shell_MissingName_ShowsUsage)
+    {
+        auto [out, err, code] = RunWsladiag(L"shell");
+        VERIFY_ARE_NOT_EQUAL(0, code);
+        VERIFY_ARE_EQUAL(L"", out);
+
+        VERIFY_IS_TRUE(err.find(L"Usage:") != std::wstring::npos);
+        VERIFY_IS_TRUE(err.find(L"wsladiag shell <SessionName> [--verbose]") != std::wstring::npos);
+    }
+
+    TEST_METHOD(Shell_InvalidSessionName_Verbose)
+    {
+        auto [out, err, code] = RunWsladiag(L"shell DefinitelyNotARealSession --verbose");
+        VERIFY_ARE_NOT_EQUAL(0, code);
+
+        VERIFY_IS_TRUE(out.find(L"[diag] shell='DefinitelyNotARealSession'") != std::wstring::npos);
+        VERIFY_IS_TRUE(err.find(L"Session not found: 'DefinitelyNotARealSession'") != std::wstring::npos);
+    }
+
+    TEST_METHOD(UnknownCommand_ShowsUsage)
+    {
+        auto [out, err, code] = RunWsladiag(L"blah");
+        VERIFY_ARE_NOT_EQUAL(0, code);
+        VERIFY_ARE_EQUAL(L"", out);
+
+        VERIFY_IS_TRUE(err.find(L"Unknown command: 'blah'") != std::wstring::npos);
+        VERIFY_IS_TRUE(err.find(L"Usage:") != std::wstring::npos);
+    }
+
+    TEST_METHOD(EmptyCommand_ShowsUsage)
+    {
+        auto [out, err, code] = RunWsladiag(L"");
+        VERIFY_ARE_EQUAL(0, code);
+        VERIFY_ARE_EQUAL(L"", out);
+
+        VERIFY_IS_TRUE(err.find(L"Usage:") != std::wstring::npos);
+        VERIFY_IS_TRUE(err.find(L"wsladiag list") != std::wstring::npos);
+        VERIFY_IS_TRUE(err.find(L"wsladiag shell <SessionName> [--verbose]") != std::wstring::npos);
+    }
+
+    TEST_METHOD(Shell_InvalidSessionName_Silent)
+    {
+        auto [out, err, code] = RunWsladiag(L"shell DefinitelyNotARealSession");
+        VERIFY_ARE_NOT_EQUAL(0, code);
+        VERIFY_ARE_EQUAL(L"", out);
+
+        VERIFY_IS_TRUE(err.find(L"Session not found: 'DefinitelyNotARealSession'") != std::wstring::npos);
+    }
+
+    TEST_METHOD(Help_ShortFlag_ShowsUsage)
+    {
+        auto [out, err, code] = RunWsladiag(L"-h");
+        VERIFY_ARE_EQUAL(0, code);
+        VERIFY_ARE_EQUAL(L"", out);
+
+        VERIFY_IS_TRUE(err.find(L"Usage:") != std::wstring::npos);
+    }
+};
+} // namespace WsladiagTests
\ No newline at end of file

From 2dd650aa6a0a74183e791c2bece5e51c8312ce75 Mon Sep 17 00:00:00 2001
From: Beena352 <beenachauhan@microsoft.com>
Date: Tue, 23 Dec 2025 15:05:12 -0800
Subject: [PATCH 06/11]  test/windows/WsladiagTests.cpp

---
 test/windows/WsladiagTests.cpp | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/test/windows/WsladiagTests.cpp b/test/windows/WsladiagTests.cpp
index bd3f6801f..5628798ee 100644
--- a/test/windows/WsladiagTests.cpp
+++ b/test/windows/WsladiagTests.cpp
@@ -130,5 +130,23 @@ class WsladiagTests
 
         VERIFY_IS_TRUE(err.find(L"Usage:") != std::wstring::npos);
     }
+
+    TEST_METHOD(Help_ShortAndLongFlags_Match)
+    {
+        auto [outH, errH, codeH] = RunWsladiag(L"-h");
+        auto [outLong, errLong, codeLong] = RunWsladiag(L"--help");
+
+        VERIFY_ARE_EQUAL(0, codeH);
+        VERIFY_ARE_EQUAL(0, codeLong);
+
+        VERIFY_ARE_EQUAL(L"", outH);
+        VERIFY_ARE_EQUAL(L"", outLong);
+
+        VERIFY_ARE_EQUAL(errH, errLong);
+
+        VERIFY_IS_TRUE(errH.find(L"Usage:") != std::wstring::npos);
+        VERIFY_IS_TRUE(errH.find(L"wsladiag list") != std::wstring::npos);
+        VERIFY_IS_TRUE(errH.find(L"wsladiag shell <SessionName> [--verbose]") != std::wstring::npos);
+    }
 };
 } // namespace WsladiagTests
\ No newline at end of file

From 7f961c8ad4e1fbbb7c7777353830b79294363632 Mon Sep 17 00:00:00 2001
From: Beena352 <beenachauhan@microsoft.com>
Date: Tue, 23 Dec 2025 15:34:57 -0800
Subject: [PATCH 07/11] Remove duplicate help flag test

---
 test/windows/WsladiagTests.cpp | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/test/windows/WsladiagTests.cpp b/test/windows/WsladiagTests.cpp
index 5628798ee..e6287d1f1 100644
--- a/test/windows/WsladiagTests.cpp
+++ b/test/windows/WsladiagTests.cpp
@@ -122,15 +122,6 @@ class WsladiagTests
         VERIFY_IS_TRUE(err.find(L"Session not found: 'DefinitelyNotARealSession'") != std::wstring::npos);
     }
 
-    TEST_METHOD(Help_ShortFlag_ShowsUsage)
-    {
-        auto [out, err, code] = RunWsladiag(L"-h");
-        VERIFY_ARE_EQUAL(0, code);
-        VERIFY_ARE_EQUAL(L"", out);
-
-        VERIFY_IS_TRUE(err.find(L"Usage:") != std::wstring::npos);
-    }
-
     TEST_METHOD(Help_ShortAndLongFlags_Match)
     {
         auto [outH, errH, codeH] = RunWsladiag(L"-h");

From 4fb58e49690285526c1e6c68029f76e687f8c422 Mon Sep 17 00:00:00 2001
From: Beena352 <beenachauhan@microsoft.com>
Date: Thu, 25 Dec 2025 13:05:29 -0800
Subject: [PATCH 08/11] Use GetMsiPackagePath() for wsladiag.exe location
 lookup

---
 test/windows/WSLATests.cpp     |   2 +-
 test/windows/WsladiagTests.cpp | 139 +++++++++++++++++++--------------
 2 files changed, 82 insertions(+), 59 deletions(-)

diff --git a/test/windows/WSLATests.cpp b/test/windows/WSLATests.cpp
index 38dfbc395..37591cb0e 100644
--- a/test/windows/WSLATests.cpp
+++ b/test/windows/WSLATests.cpp
@@ -170,7 +170,7 @@ class WSLATests
         if (result.Code != expectedResult)
         {
             LogError(
-                "Comman didn't return expected code (%i). ExitCode: %i, Stdout: '%hs', Stderr: '%hs'",
+                "Command didn't return expected code (%i). ExitCode: %i, Stdout: '%hs', Stderr: '%hs'",
                 expectedResult,
                 result.Code,
                 result.Output[1].c_str(),
diff --git a/test/windows/WsladiagTests.cpp b/test/windows/WsladiagTests.cpp
index e6287d1f1..15c9d1cc8 100644
--- a/test/windows/WsladiagTests.cpp
+++ b/test/windows/WsladiagTests.cpp
@@ -21,6 +21,7 @@ class WsladiagTests
 {
     WSL_TEST_CLASS(WsladiagTests)
 
+    // Initialize the tests
     TEST_CLASS_SETUP(TestClassSetup)
     {
         VERIFY_ARE_EQUAL(LxsstuInitialize(FALSE), TRUE);
@@ -32,66 +33,51 @@ class WsladiagTests
         LxsstuUninitialize(FALSE);
         return true;
     }
-
-    static std::wstring BuildWsladiagCmd(const std::wstring& args)
-    {
-        const auto exePath = wsl::windows::common::wslutil::GetBasePath() / L"wsladiag.exe";
-        VERIFY_IS_TRUE(std::filesystem::exists(exePath));
-
-        const auto exe = exePath.wstring();
-        return args.empty() ? std::format(L"\"{}\"", exe) : std::format(L"\"{}\" {}", exe, args);
-    }
-
-    static std::tuple<std::wstring, std::wstring, int> RunWsladiag(const std::wstring& args)
-    {
-        auto cmd = BuildWsladiagCmd(args);
-        return LxsstuLaunchCommandAndCaptureOutputWithResult(cmd.data());
-    }
-
+    // Test that wsladiag list command shows either sessions or "no sessions" message
     TEST_METHOD(List_ShowsSessionsOrNoSessions)
     {
         auto [out, err, code] = RunWsladiag(L"list");
         VERIFY_ARE_EQUAL(0, code);
         VERIFY_ARE_EQUAL(L"", err);
 
-        const bool noSessions = (out.find(L"No WSLA sessions found.") != std::wstring::npos);
-
-        const bool hasTable = (out.find(L"WSLA session") != std::wstring::npos) && (out.find(L"ID") != std::wstring::npos) &&
-                              (out.find(L"Display Name") != std::wstring::npos);
-
-        VERIFY_IS_TRUE(noSessions || hasTable);
+        ValidateListOutput(out);
     }
 
+    // Test that wsladiag --help shows usage information
     TEST_METHOD(Help_ShowsUsage)
     {
         auto [out, err, code] = RunWsladiag(L"--help");
         VERIFY_ARE_EQUAL(0, code);
         VERIFY_ARE_EQUAL(L"", out);
-
-        VERIFY_IS_TRUE(err.find(L"Usage:") != std::wstring::npos);
-        VERIFY_IS_TRUE(err.find(L"wsladiag list") != std::wstring::npos);
-        VERIFY_IS_TRUE(err.find(L"wsladiag shell <SessionName> [--verbose]") != std::wstring::npos);
+        ValidateUsage(err);
     }
 
-    TEST_METHOD(Shell_MissingName_ShowsUsage)
+    // Test that wsladiag with no arguments shows usage information
+    TEST_METHOD(EmptyCommand_ShowsUsage)
     {
-        auto [out, err, code] = RunWsladiag(L"shell");
-        VERIFY_ARE_NOT_EQUAL(0, code);
+        auto [out, err, code] = RunWsladiag(L"");
+        VERIFY_ARE_EQUAL(0, code);
         VERIFY_ARE_EQUAL(L"", out);
-
-        VERIFY_IS_TRUE(err.find(L"Usage:") != std::wstring::npos);
-        VERIFY_IS_TRUE(err.find(L"wsladiag shell <SessionName> [--verbose]") != std::wstring::npos);
+        ValidateUsage(err);
     }
 
-    TEST_METHOD(Shell_InvalidSessionName_Verbose)
+    // Test that -h and --help flags produce identical output
+    TEST_METHOD(Help_ShortAndLongFlags_Match)
     {
-        auto [out, err, code] = RunWsladiag(L"shell DefinitelyNotARealSession --verbose");
-        VERIFY_ARE_NOT_EQUAL(0, code);
+        auto [outH, errH, codeH] = RunWsladiag(L"-h");
+        auto [outLong, errLong, codeLong] = RunWsladiag(L"--help");
+
+        VERIFY_ARE_EQUAL(0, codeH);
+        VERIFY_ARE_EQUAL(0, codeLong);
 
-        VERIFY_IS_TRUE(out.find(L"[diag] shell='DefinitelyNotARealSession'") != std::wstring::npos);
-        VERIFY_IS_TRUE(err.find(L"Session not found: 'DefinitelyNotARealSession'") != std::wstring::npos);
+        VERIFY_ARE_EQUAL(L"", outH);
+        VERIFY_ARE_EQUAL(L"", outLong);
+
+        VERIFY_ARE_EQUAL(errH, errLong);
+        ValidateUsage(errH);
     }
 
+    // Test that unknown commands show error message and usage
     TEST_METHOD(UnknownCommand_ShowsUsage)
     {
         auto [out, err, code] = RunWsladiag(L"blah");
@@ -99,45 +85,82 @@ class WsladiagTests
         VERIFY_ARE_EQUAL(L"", out);
 
         VERIFY_IS_TRUE(err.find(L"Unknown command: 'blah'") != std::wstring::npos);
-        VERIFY_IS_TRUE(err.find(L"Usage:") != std::wstring::npos);
+        ValidateUsage(err);
     }
 
-    TEST_METHOD(EmptyCommand_ShowsUsage)
+    // Test that shell command without session name shows usage
+    TEST_METHOD(Shell_MissingName_ShowsUsage)
     {
-        auto [out, err, code] = RunWsladiag(L"");
-        VERIFY_ARE_EQUAL(0, code);
+        auto [out, err, code] = RunWsladiag(L"shell");
+        VERIFY_ARE_NOT_EQUAL(0, code);
         VERIFY_ARE_EQUAL(L"", out);
-
-        VERIFY_IS_TRUE(err.find(L"Usage:") != std::wstring::npos);
-        VERIFY_IS_TRUE(err.find(L"wsladiag list") != std::wstring::npos);
-        VERIFY_IS_TRUE(err.find(L"wsladiag shell <SessionName> [--verbose]") != std::wstring::npos);
+        ValidateUsage(err);
     }
 
+    // Test shell command with invalid session name (silent mode)
     TEST_METHOD(Shell_InvalidSessionName_Silent)
     {
-        auto [out, err, code] = RunWsladiag(L"shell DefinitelyNotARealSession");
+        const std::wstring name = L"DefinitelyNotARealSession";
+        auto [out, err, code] = RunWsladiag(std::format(L"shell {}", name));
         VERIFY_ARE_NOT_EQUAL(0, code);
         VERIFY_ARE_EQUAL(L"", out);
 
-        VERIFY_IS_TRUE(err.find(L"Session not found: 'DefinitelyNotARealSession'") != std::wstring::npos);
+        ValidateSessionNotFound(err, name);
     }
 
-    TEST_METHOD(Help_ShortAndLongFlags_Match)
+    // Test shell command with invalid session name (verbose mode)
+    TEST_METHOD(Shell_InvalidSessionName_Verbose)
     {
-        auto [outH, errH, codeH] = RunWsladiag(L"-h");
-        auto [outLong, errLong, codeLong] = RunWsladiag(L"--help");
+        const std::wstring name = L"DefinitelyNotARealSession";
+        auto [out, err, code] = RunWsladiag(std::format(L"shell {} --verbose", name));
+        VERIFY_ARE_NOT_EQUAL(0, code);
 
-        VERIFY_ARE_EQUAL(0, codeH);
-        VERIFY_ARE_EQUAL(0, codeLong);
+        VERIFY_IS_TRUE(out.find(std::format(L"[diag] shell='{}'", name)) != std::wstring::npos);
+        ValidateSessionNotFound(err, name);
+    }
 
-        VERIFY_ARE_EQUAL(L"", outH);
-        VERIFY_ARE_EQUAL(L"", outLong);
+    // Build command line for wsladiag.exe with given arguments
+    static std::wstring BuildWsladiagCmd(const std::wstring& args)
+    {
+        const auto msiPathOpt = wsl::windows::common::wslutil::GetMsiPackagePath();
+        VERIFY_IS_TRUE(msiPathOpt.has_value());
 
-        VERIFY_ARE_EQUAL(errH, errLong);
+        const auto exePath = std::filesystem::path(*msiPathOpt) / L"wsladiag.exe";
+        const auto exe = exePath.wstring();
 
-        VERIFY_IS_TRUE(errH.find(L"Usage:") != std::wstring::npos);
-        VERIFY_IS_TRUE(errH.find(L"wsladiag list") != std::wstring::npos);
-        VERIFY_IS_TRUE(errH.find(L"wsladiag shell <SessionName> [--verbose]") != std::wstring::npos);
+        return args.empty() ? std::format(L"\"{}\"", exe) : std::format(L"\"{}\" {}", exe, args);
+    }
+
+    // Execute wsladiag with given arguments and return output, error, and exit code
+    static std::tuple<std::wstring, std::wstring, int> RunWsladiag(const std::wstring& args)
+    {
+        auto cmd = BuildWsladiagCmd(args);
+        return LxsstuLaunchCommandAndCaptureOutputWithResult(cmd.data());
+    }
+
+    // Validate that list command output shows either no sessions message or session table
+    static void ValidateListOutput(const std::wstring& out)
+    {
+        const bool noSessions = out.find(L"No WSLA sessions found.") != std::wstring::npos;
+
+        const bool hasTable = out.find(L"WSLA session") != std::wstring::npos && out.find(L"ID") != std::wstring::npos &&
+                              out.find(L"Creator PID") != std::wstring::npos && out.find(L"Display Name") != std::wstring::npos;
+
+        VERIFY_IS_TRUE(noSessions || hasTable);
+    }
+
+    // Validate that usage information contains expected command descriptions
+    static void ValidateUsage(const std::wstring& err)
+    {
+        VERIFY_IS_TRUE(err.find(L"Usage:") != std::wstring::npos);
+        VERIFY_IS_TRUE(err.find(L"wsladiag list") != std::wstring::npos);
+        VERIFY_IS_TRUE(err.find(L"wsladiag shell <SessionName> [--verbose]") != std::wstring::npos);
+    }
+
+    // Validate that session not found error contains the expected session name
+    static void ValidateSessionNotFound(const std::wstring& err, const std::wstring& name)
+    {
+        VERIFY_IS_TRUE(err.find(std::format(L"Session not found: '{}'", name)) != std::wstring::npos);
     }
 };
 } // namespace WsladiagTests
\ No newline at end of file

From b7d8bb518dba47d45dc7e96517902794b6bace2d Mon Sep 17 00:00:00 2001
From: Beena352 <beenachauhan@microsoft.com>
Date: Fri, 26 Dec 2025 15:27:58 -0800
Subject: [PATCH 09/11] Refine wsladiag test assertions

---
 src/windows/common/ExecutionContext.h |  1 +
 src/windows/common/wslutil.cpp        |  3 +-
 src/windows/wsladiag/wsladiag.cpp     | 34 ++++++++---
 test/windows/WsladiagTests.cpp        | 81 ++++++++++++---------------
 4 files changed, 67 insertions(+), 52 deletions(-)

diff --git a/src/windows/common/ExecutionContext.h b/src/windows/common/ExecutionContext.h
index 0fe473d9f..9c12ff644 100644
--- a/src/windows/common/ExecutionContext.h
+++ b/src/windows/common/ExecutionContext.h
@@ -66,6 +66,7 @@ enum Context : ULONGLONG
     UpdatePackage = 0x10000000000,
     QueryLatestGitHubRelease = 0x20000000000,
     VerifyChecksum = 0x40000000000,
+    WslaDiag = 0x80000000000,
 };
 
 DEFINE_ENUM_FLAG_OPERATORS(Context)
diff --git a/src/windows/common/wslutil.cpp b/src/windows/common/wslutil.cpp
index 8c4ad8ac3..544c6fce5 100644
--- a/src/windows/common/wslutil.cpp
+++ b/src/windows/common/wslutil.cpp
@@ -190,7 +190,8 @@ static const std::map<Context, LPCWSTR> g_contextStrings{
     X(HNS),
     X(ReadDistroConfig),
     X(MoveDistro),
-    X(VerifyChecksum)};
+    X(VerifyChecksum),
+    X(WslaDiag)};
 
 #undef X
 
diff --git a/src/windows/wsladiag/wsladiag.cpp b/src/windows/wsladiag/wsladiag.cpp
index 36fa5db7e..ace8c8296 100644
--- a/src/windows/wsladiag/wsladiag.cpp
+++ b/src/windows/wsladiag/wsladiag.cpp
@@ -18,11 +18,14 @@ Module Name:
 #include "wslaservice.h"
 #include "WslSecurity.h"
 #include "WSLAProcessLauncher.h"
+#include "ExecutionContext.h"
 #include <thread>
 #include <format>
 
 using namespace wsl::shared;
 namespace wslutil = wsl::windows::common::wslutil;
+using wsl::windows::common::Context;
+using wsl::windows::common::ExecutionContext;
 using wsl::windows::common::WSLAProcessLauncher;
 
 // Adding a helper to factor error handling between all the arguments.
@@ -319,18 +322,35 @@ int wsladiag_main(std::wstring_view commandLine)
 
 int wmain(int, wchar_t**)
 {
+    wsl::windows::common::EnableContextualizedErrors(false);
+
+    std::optional<ExecutionContext> context;
+    int exitCode = 1;
+    HRESULT result = S_OK;
+
     try
     {
-        return wsladiag_main(GetCommandLineW());
+        context.emplace(Context::WslaDiag);
+        exitCode = wsladiag_main(GetCommandLineW());
     }
     catch (...)
     {
-        const auto hr = wil::ResultFromCaughtException();
-        if (hr == E_INVALIDARG)
+        result = wil::ResultFromCaughtException();
+    }
+
+    if (FAILED(result))
+    {
+        if (context.has_value() && context->ReportedError().has_value())
         {
-            PrintUsage();
-            return 1;
+            auto strings = wsl::windows::common::wslutil::ErrorToString(context->ReportedError().value());
+            wslutil::PrintMessage(strings.Message, stderr);
+        }
+        else
+        {
+            // Fallback for errors without context
+            wslutil::PrintMessage(wslutil::GetErrorString(result), stderr);
         }
-        return ReportError(L"wsladiag failed", hr);
     }
-}
+
+    return exitCode;
+}
\ No newline at end of file
diff --git a/test/windows/WsladiagTests.cpp b/test/windows/WsladiagTests.cpp
index 15c9d1cc8..f508d24d1 100644
--- a/test/windows/WsladiagTests.cpp
+++ b/test/windows/WsladiagTests.cpp
@@ -21,18 +21,6 @@ class WsladiagTests
 {
     WSL_TEST_CLASS(WsladiagTests)
 
-    // Initialize the tests
-    TEST_CLASS_SETUP(TestClassSetup)
-    {
-        VERIFY_ARE_EQUAL(LxsstuInitialize(FALSE), TRUE);
-        return true;
-    }
-
-    TEST_CLASS_CLEANUP(TestClassCleanup)
-    {
-        LxsstuUninitialize(FALSE);
-        return true;
-    }
     // Test that wsladiag list command shows either sessions or "no sessions" message
     TEST_METHOD(List_ShowsSessionsOrNoSessions)
     {
@@ -46,19 +34,13 @@ class WsladiagTests
     // Test that wsladiag --help shows usage information
     TEST_METHOD(Help_ShowsUsage)
     {
-        auto [out, err, code] = RunWsladiag(L"--help");
-        VERIFY_ARE_EQUAL(0, code);
-        VERIFY_ARE_EQUAL(L"", out);
-        ValidateUsage(err);
+        ValidateWslaDiagOutput(L"--help", 0, L"Usage:");
     }
 
     // Test that wsladiag with no arguments shows usage information
     TEST_METHOD(EmptyCommand_ShowsUsage)
     {
-        auto [out, err, code] = RunWsladiag(L"");
-        VERIFY_ARE_EQUAL(0, code);
-        VERIFY_ARE_EQUAL(L"", out);
-        ValidateUsage(err);
+        ValidateWslaDiagOutput(L"", 0, L"Usage:");
     }
 
     // Test that -h and --help flags produce identical output
@@ -78,34 +60,26 @@ class WsladiagTests
     }
 
     // Test that unknown commands show error message and usage
-    TEST_METHOD(UnknownCommand_ShowsUsage)
+    TEST_METHOD(UnknownCommand_ShowsError)
     {
         auto [out, err, code] = RunWsladiag(L"blah");
-        VERIFY_ARE_NOT_EQUAL(0, code);
-        VERIFY_ARE_EQUAL(L"", out);
+        VERIFY_ARE_EQUAL(1, code);
 
-        VERIFY_IS_TRUE(err.find(L"Unknown command: 'blah'") != std::wstring::npos);
-        ValidateUsage(err);
+        const std::wstring combined = out + err;
+        VERIFY_IS_TRUE(combined.find(L"Unknown command: 'blah'") != std::wstring::npos);
+        VERIFY_IS_TRUE(combined.find(L"Usage:") != std::wstring::npos);
     }
 
-    // Test that shell command without session name shows usage
-    TEST_METHOD(Shell_MissingName_ShowsUsage)
+    // Test that shell command without session name shows error
+    TEST_METHOD(Shell_MissingName_ShowsError)
     {
-        auto [out, err, code] = RunWsladiag(L"shell");
-        VERIFY_ARE_NOT_EQUAL(0, code);
-        VERIFY_ARE_EQUAL(L"", out);
-        ValidateUsage(err);
+        ValidateWslaDiagFailsWith(L"shell", L"wsladiag shell <SessionName> [--verbose]");
     }
 
     // Test shell command with invalid session name (silent mode)
     TEST_METHOD(Shell_InvalidSessionName_Silent)
     {
-        const std::wstring name = L"DefinitelyNotARealSession";
-        auto [out, err, code] = RunWsladiag(std::format(L"shell {}", name));
-        VERIFY_ARE_NOT_EQUAL(0, code);
-        VERIFY_ARE_EQUAL(L"", out);
-
-        ValidateSessionNotFound(err, name);
+        ValidateWslaDiagFailsWith(L"shell DefinitelyNotARealSession", L"Session not found: 'DefinitelyNotARealSession'");
     }
 
     // Test shell command with invalid session name (verbose mode)
@@ -116,7 +90,7 @@ class WsladiagTests
         VERIFY_ARE_NOT_EQUAL(0, code);
 
         VERIFY_IS_TRUE(out.find(std::format(L"[diag] shell='{}'", name)) != std::wstring::npos);
-        ValidateSessionNotFound(err, name);
+        VERIFY_IS_TRUE(err.find(L"Session not found") != std::wstring::npos);
     }
 
     // Build command line for wsladiag.exe with given arguments
@@ -138,6 +112,31 @@ class WsladiagTests
         return LxsstuLaunchCommandAndCaptureOutputWithResult(cmd.data());
     }
 
+    static void ValidateWslaDiagOutput(const std::wstring& cmd, const std::wstring& expectedSubstring)
+    {
+        auto [out, err, code] = RunWsladiag(cmd);
+        const std::wstring combined = out + err;
+        VERIFY_IS_TRUE(combined.find(expectedSubstring) != std::wstring::npos);
+    }
+
+    static void ValidateWslaDiagOutput(const std::wstring& cmd, int expectedExitCode, const std::wstring& expectedSubstring)
+    {
+        auto [out, err, code] = RunWsladiag(cmd);
+        VERIFY_ARE_EQUAL(expectedExitCode, code);
+
+        const std::wstring combined = out + err;
+        VERIFY_IS_TRUE(combined.find(expectedSubstring) != std::wstring::npos);
+    }
+
+    static void ValidateWslaDiagFailsWith(const std::wstring& cmd, const std::wstring& expectedSubstring)
+    {
+        auto [out, err, code] = RunWsladiag(cmd);
+        VERIFY_ARE_NOT_EQUAL(0, code);
+
+        const std::wstring combined = out + err;
+        VERIFY_IS_TRUE(combined.find(expectedSubstring) != std::wstring::npos);
+    }
+
     // Validate that list command output shows either no sessions message or session table
     static void ValidateListOutput(const std::wstring& out)
     {
@@ -156,11 +155,5 @@ class WsladiagTests
         VERIFY_IS_TRUE(err.find(L"wsladiag list") != std::wstring::npos);
         VERIFY_IS_TRUE(err.find(L"wsladiag shell <SessionName> [--verbose]") != std::wstring::npos);
     }
-
-    // Validate that session not found error contains the expected session name
-    static void ValidateSessionNotFound(const std::wstring& err, const std::wstring& name)
-    {
-        VERIFY_IS_TRUE(err.find(std::format(L"Session not found: '{}'", name)) != std::wstring::npos);
-    }
 };
 } // namespace WsladiagTests
\ No newline at end of file

From c5a97ae2ee05005288e7692e00e1d14f43b22aa3 Mon Sep 17 00:00:00 2001
From: Beena352 <beenachauhan@microsoft.com>
Date: Fri, 26 Dec 2025 16:18:03 -0800
Subject: [PATCH 10/11] Fix WsladiagTests Shell_MissingName_ShowsError
 validation

---
 test/windows/WsladiagTests.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/windows/WsladiagTests.cpp b/test/windows/WsladiagTests.cpp
index f508d24d1..96389eceb 100644
--- a/test/windows/WsladiagTests.cpp
+++ b/test/windows/WsladiagTests.cpp
@@ -73,7 +73,7 @@ class WsladiagTests
     // Test that shell command without session name shows error
     TEST_METHOD(Shell_MissingName_ShowsError)
     {
-        ValidateWslaDiagFailsWith(L"shell", L"wsladiag shell <SessionName> [--verbose]");
+        ValidateWslaDiagFailsWith(L"shell", L"The parameter is incorrect.");
     }
 
     // Test shell command with invalid session name (silent mode)

From eb7463b31d83e6fcbe17270e19f69dc32b1147d4 Mon Sep 17 00:00:00 2001
From: Beena352 <beenachauhan@microsoft.com>
Date: Sat, 3 Jan 2026 17:43:45 -0800
Subject: [PATCH 11/11] tests: validate wsladiag stdout/stderr

---
 src/windows/wsladiag/wsladiag.cpp |  9 +++---
 test/windows/WsladiagTests.cpp    | 51 ++++++++++++-------------------
 2 files changed, 24 insertions(+), 36 deletions(-)

diff --git a/src/windows/wsladiag/wsladiag.cpp b/src/windows/wsladiag/wsladiag.cpp
index ace8c8296..0cec3a59c 100644
--- a/src/windows/wsladiag/wsladiag.cpp
+++ b/src/windows/wsladiag/wsladiag.cpp
@@ -324,13 +324,12 @@ int wmain(int, wchar_t**)
 {
     wsl::windows::common::EnableContextualizedErrors(false);
 
-    std::optional<ExecutionContext> context;
+    ExecutionContext context{Context::WslaDiag};
     int exitCode = 1;
     HRESULT result = S_OK;
 
     try
     {
-        context.emplace(Context::WslaDiag);
         exitCode = wsladiag_main(GetCommandLineW());
     }
     catch (...)
@@ -340,10 +339,10 @@ int wmain(int, wchar_t**)
 
     if (FAILED(result))
     {
-        if (context.has_value() && context->ReportedError().has_value())
+        if (auto reported = context.ReportedError())
         {
-            auto strings = wsl::windows::common::wslutil::ErrorToString(context->ReportedError().value());
-            wslutil::PrintMessage(strings.Message, stderr);
+            auto strings = wsl::windows::common::wslutil::ErrorToString(*reported);
+            wslutil::PrintMessage(strings.Message.empty() ? strings.Code : strings.Message, stderr);
         }
         else
         {
diff --git a/test/windows/WsladiagTests.cpp b/test/windows/WsladiagTests.cpp
index 96389eceb..50bed9fc9 100644
--- a/test/windows/WsladiagTests.cpp
+++ b/test/windows/WsladiagTests.cpp
@@ -16,6 +16,13 @@ Module Name:
 #include "Common.h"
 #include <format>
 
+static const std::wstring c_usageText =
+    L"wsladiag - WSLA diagnostics tool\r\n"
+    L"Usage:\r\n"
+    L"  wsladiag list\r\n"
+    L"  wsladiag shell <SessionName> [--verbose]\r\n"
+    L"  wsladiag --help\r\n";
+
 namespace WsladiagTests {
 class WsladiagTests
 {
@@ -34,13 +41,13 @@ class WsladiagTests
     // Test that wsladiag --help shows usage information
     TEST_METHOD(Help_ShowsUsage)
     {
-        ValidateWslaDiagOutput(L"--help", 0, L"Usage:");
+        ValidateWsladiagOutput(L"--help", 0, L"", c_usageText);
     }
 
     // Test that wsladiag with no arguments shows usage information
     TEST_METHOD(EmptyCommand_ShowsUsage)
     {
-        ValidateWslaDiagOutput(L"", 0, L"Usage:");
+        ValidateWsladiagOutput(L"", 0, L"", c_usageText);
     }
 
     // Test that -h and --help flags produce identical output
@@ -62,24 +69,23 @@ class WsladiagTests
     // Test that unknown commands show error message and usage
     TEST_METHOD(UnknownCommand_ShowsError)
     {
-        auto [out, err, code] = RunWsladiag(L"blah");
-        VERIFY_ARE_EQUAL(1, code);
-
-        const std::wstring combined = out + err;
-        VERIFY_IS_TRUE(combined.find(L"Unknown command: 'blah'") != std::wstring::npos);
-        VERIFY_IS_TRUE(combined.find(L"Usage:") != std::wstring::npos);
+        ValidateWsladiagOutput(L"blah", 1, L"", std::wstring(L"Unknown command: 'blah'\r\n") + c_usageText);
     }
 
     // Test that shell command without session name shows error
     TEST_METHOD(Shell_MissingName_ShowsError)
     {
-        ValidateWslaDiagFailsWith(L"shell", L"The parameter is incorrect.");
+        ValidateWsladiagOutput(L"shell", 1, L"", L"The parameter is incorrect.\r\n");
     }
 
     // Test shell command with invalid session name (silent mode)
     TEST_METHOD(Shell_InvalidSessionName_Silent)
     {
-        ValidateWslaDiagFailsWith(L"shell DefinitelyNotARealSession", L"Session not found: 'DefinitelyNotARealSession'");
+        auto [out, err, code] = RunWsladiag(L"shell DefinitelyNotARealSession");
+        VERIFY_ARE_NOT_EQUAL(0, code);
+
+        VERIFY_ARE_EQUAL(L"", out);
+        VERIFY_IS_TRUE(err.find(L"Session not found: 'DefinitelyNotARealSession'") != std::wstring::npos);
     }
 
     // Test shell command with invalid session name (verbose mode)
@@ -112,29 +118,12 @@ class WsladiagTests
         return LxsstuLaunchCommandAndCaptureOutputWithResult(cmd.data());
     }
 
-    static void ValidateWslaDiagOutput(const std::wstring& cmd, const std::wstring& expectedSubstring)
-    {
-        auto [out, err, code] = RunWsladiag(cmd);
-        const std::wstring combined = out + err;
-        VERIFY_IS_TRUE(combined.find(expectedSubstring) != std::wstring::npos);
-    }
-
-    static void ValidateWslaDiagOutput(const std::wstring& cmd, int expectedExitCode, const std::wstring& expectedSubstring)
+    static void ValidateWsladiagOutput(const std::wstring& cmd, int expectedExitCode, const std::wstring& expectedStdout, const std::wstring& expectedStderr)
     {
         auto [out, err, code] = RunWsladiag(cmd);
         VERIFY_ARE_EQUAL(expectedExitCode, code);
-
-        const std::wstring combined = out + err;
-        VERIFY_IS_TRUE(combined.find(expectedSubstring) != std::wstring::npos);
-    }
-
-    static void ValidateWslaDiagFailsWith(const std::wstring& cmd, const std::wstring& expectedSubstring)
-    {
-        auto [out, err, code] = RunWsladiag(cmd);
-        VERIFY_ARE_NOT_EQUAL(0, code);
-
-        const std::wstring combined = out + err;
-        VERIFY_IS_TRUE(combined.find(expectedSubstring) != std::wstring::npos);
+        VERIFY_ARE_EQUAL(expectedStdout, out);
+        VERIFY_ARE_EQUAL(expectedStderr, err);
     }
 
     // Validate that list command output shows either no sessions message or session table
@@ -142,7 +131,7 @@ class WsladiagTests
     {
         const bool noSessions = out.find(L"No WSLA sessions found.") != std::wstring::npos;
 
-        const bool hasTable = out.find(L"WSLA session") != std::wstring::npos && out.find(L"ID") != std::wstring::npos &&
+        const bool hasTable = out.find(L"Found") != std::wstring::npos && out.find(L"ID") != std::wstring::npos &&
                               out.find(L"Creator PID") != std::wstring::npos && out.find(L"Display Name") != std::wstring::npos;
 
         VERIFY_IS_TRUE(noSessions || hasTable);