[lldb-dap] bug-fix: lldb-dap remembered old sessionId when reused

Summary:
Context:
lldb-dap remembers the vscode session id in its lifecycle,
the new session id when reuse lldb-dap didn't get updated in lldb-dap.

This diff is aiming to fix such issue on the lldb-dap side.

Approach:
fetching `__sessionId` key from attach & launch request, and update the env value by inserting new initCommands.

Test Plan:
Added new test case to test the env value update in lldb-dap

to run the test case:
~/local/llvm-sand/build/Debug/fbcode-x86_64/toolchain/bin/lldb-dotest -p <PATH_TO_TEST>/TestDAP_launch.py
^ replace path to the local build lldb-dotest.
^ [lldb-dotest](https://www.internalfb.com/wiki/SaND/Developing_LLDB)

Attach Test cases results: P1754726175
Launch Test Cases results: P1754738053
reuseDap Test Cases results: P1754757203

Reviewers: jeffreytan, royshi, gclayton, #lldb_team

Reviewed By: jeffreytan, royshi

Subscribers: royshi

Differential Revision: https://phabricator.intern.facebook.com/D70663297

Tasks: T215634007

Author: zyn-li

lldb/packages/Python/lldbsuite/test/tools/lldb-dap/dap_server.py

@@ -697,6 +697,7 @@ def request_attach(
         sourceMap: Optional[Union[list[tuple[str, str]], dict[str, str]]] = None,
         gdbRemotePort: Optional[int] = None,
         gdbRemoteHostname: Optional[str] = None,
+        vscode_session_id=None,
     ):
         args_dict = {}
         if pid is not None:
@@ -730,6 +731,9 @@ def request_attach(
             args_dict["gdb-remote-port"] = gdbRemotePort
         if gdbRemoteHostname is not None:
             args_dict["gdb-remote-hostname"] = gdbRemoteHostname
+        if vscode_session_id: 
+            args_dict["__sessionId"] = vscode_session_id
+            
         command_dict = {"command": "attach", "type": "request", "arguments": args_dict}
         return self.send_recv(command_dict)
 
@@ -939,6 +943,7 @@ def request_launch(
         commandEscapePrefix: Optional[str] = None,
         customFrameFormat: Optional[str] = None,
         customThreadFormat: Optional[str] = None,
+        vscode_session_id=None,
     ):
         args_dict = {"program": program}
         if args:
@@ -980,6 +985,8 @@ def request_launch(
             args_dict["customFrameFormat"] = customFrameFormat
         if customThreadFormat:
             args_dict["customThreadFormat"] = customThreadFormat
+        if vscode_session_id: 
+            args_dict["__sessionId"] = vscode_session_id
 
         args_dict["disableASLR"] = disableASLR
         args_dict["enableAutoVariableSummaries"] = enableAutoVariableSummaries

lldb/test/API/tools/lldb-dap/attach/TestDAP_attach.py

@@ -227,3 +227,45 @@ def test_terminate_commands(self):
             pattern=terminateCommands[0],
         )
         self.verify_commands("terminateCommands", output, terminateCommands)
+    
+    def test_session_id_update(self):
+        self.build_and_create_debug_adapter()
+        program = self.getBuildArtifact("a.out")
+        
+        def spawn(program):
+            process = subprocess.Popen(
+                [program], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+            )
+        self.spawn_thread = threading.Thread(
+            target=spawn,
+            args=(program,),
+        )
+        self.spawn_thread.start()
+        postRunCommands = ["script print('Actual_Session_ID: ' + str(os.getenv('VSCODE_DEBUG_SESSION_ID')))"]
+        self.attach(program=program, vscode_session_id="test_session_id", postRunCommands=postRunCommands)
+        output = self.get_console() 
+        lines = filter(lambda x: 'Actual_Session_ID' in x, output.splitlines())  
+        self.assertTrue(
+            any('test_session_id' in l for l in lines), "expect session id in console output"
+        )
+    
+    def test_session_id_update_empty(self): 
+        self.build_and_create_debug_adapter()
+        program = self.getBuildArtifact("a.out")
+        
+        def spawn(program):
+            process = subprocess.Popen(
+                [program], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+            )
+        self.spawn_thread = threading.Thread(
+            target=spawn,
+            args=(program,),
+        )
+        self.spawn_thread.start()
+        postRunCommands = ["script print('Actual_Session_ID: ' + str(os.getenv('VSCODE_DEBUG_SESSION_ID', 'None')))"]
+        self.attach(program=program, postRunCommands=postRunCommands)
+        output = self.get_console() 
+        lines = filter(lambda x: 'Actual_Session_ID' in x, output.splitlines())
+        self.assertTrue(
+            any(l == "Actual_Session_ID: None" in l for l in lines), "expect session id in console output"
+        )

lldb/test/API/tools/lldb-dap/launch/TestDAP_launch.py

@@ -572,3 +572,28 @@ def test_version(self):
             version_string.splitlines(),
             "version string does not match",
         )
+
+    def test_session_id_update(self):
+        program = self.getBuildArtifact("a.out")
+        postRunCommands = ["script print('Actual_Session_ID: ' + str(os.getenv('VSCODE_DEBUG_SESSION_ID')))"]
+        self.build_and_launch(
+            program, 
+            vscode_session_id="test_session_id", 
+            postRunCommands=postRunCommands,
+        )  
+        output = self.get_console() 
+        self.continue_to_exit()
+        lines = filter(lambda x: 'Actual_Session_ID' in x, output.splitlines())  
+        self.assertTrue(
+            any("test_session_id" in l for l in lines), "expect session id in console output"
+        )
+    
+    def test_session_id_update_empty(self): 
+        program = self.getBuildArtifact("a.out")
+        self.build_and_launch(program)  
+        output = self.get_console() 
+        self.continue_to_exit()
+        self.assertTrue(
+            all("VSCODE_DEBUG_SESSION_ID" not in l for l in output.splitlines()), 
+            "expect NO session id update command"
+        )

lldb/test/API/tools/lldb-dap/reuseDAP/TestDAP_reuseAdapter.py

@@ -11,7 +11,7 @@
 import os
 
 
-class TestDAP_reuseAdpater(lldbdap_testcase.DAPTestCaseBase):
+class TestDAP_reuseAdapter(lldbdap_testcase.DAPTestCaseBase):
     @skipIfWindows
     def test_basic_reuse(self):
         """
@@ -20,7 +20,7 @@ def test_basic_reuse(self):
         program = self.getBuildArtifact("a.out")
 
         # Keep lldb-dap alive for 10 minutes.
-        dapKeepAliveTimeInMS = 10 * 1000 * 60 
+        dapKeepAliveTimeInMS = 10 * 1000 * 60
         self.build_and_launch(program, disconnectAutomatically=False, keepAliveTimeout=dapKeepAliveTimeInMS)
 
         source = "main.cpp"
@@ -32,7 +32,7 @@ def test_basic_reuse(self):
         # Second debug session by reusing lldb-dap.
         self.create_debug_adapter(reuseDapServer=True)
         self.launch(program)
-        
+
         breakpoint2_line = line_number(source, "// breakpoint 2")
         breakpoint_ids = self.set_source_breakpoints(source, [breakpoint2_line])
         self.continue_to_breakpoints(breakpoint_ids)
@@ -68,3 +68,45 @@ def test_exception_breakopints(self):
         self.assertTrue(response["success"])
         self.continue_to_exception_breakpoint("C++ Throw")
         self.dap_server.request_disconnect()
+
+    def test_session_id_update_when_reuse(self):
+        """
+        Test if vscode session id being updated when lldb-dap reused
+        """
+        # Add set up command to print out the session id, with a unique identifier as prefix
+        program = self.getBuildArtifact("a.out")
+        postRunCommands = ["script print('Actual_Session_ID: ' + str(os.getenv('VSCODE_DEBUG_SESSION_ID')))"]
+
+        # Keep lldb-dap alive for 10 minutes.
+        dapKeepAliveTimeInMS = 10 * 1000 * 60
+        self.build_and_launch(
+            program,
+            vscode_session_id="test_session_id",
+            postRunCommands=postRunCommands,
+            disconnectAutomatically=False,
+            keepAliveTimeout=dapKeepAliveTimeInMS
+        )
+
+        # Validate the session id in the inital launch of lldb-dap
+        output = self.get_console()
+        lines = filter(lambda x: 'Actual_Session_ID' in x, output.splitlines())
+        self.assertTrue(
+            any("test_session_id" in l for l in lines), "expect session id in console output"
+        )
+        self.dap_server.request_disconnect()
+
+        # Second debug session by reusing lldb-dap.
+        self.create_debug_adapter(reuseDapServer=True)
+        self.launch(
+            program,
+            vscode_session_id="NEW_session_id",
+            postRunCommands=postRunCommands
+        )
+
+        # Validate the updated session id in the reused lldb-dap
+        output = self.get_console()
+        lines = filter(lambda x: 'Actual_Session_ID' in x, output.splitlines())
+        self.assertTrue(
+            any("NEW_session_id" in l for l in lines), "expect new session id in console output"
+        )
+        self.dap_server.request_disconnect()

lldb/tools/lldb-dap/Handler/AttachRequestHandler.cpp

@@ -48,6 +48,12 @@ Error AttachRequestHandler::Run(const AttachRequestArguments &args) const {
   if (!args.coreFile.empty())
     dap.stop_at_entry = true;
 
+  // When lldb-dap is reused, the ID of each session is passed in by a request
+  // argument, rather than through the environment variable. Set it into the
+  // environment variable, so that the downstream code can pick it up from there.
+  // See T215634007 for more details.
+  UpdateVSCodeSessionID(args, dap);
+
   PrintWelcomeMessage();
 
   // This is a hack for loading DWARF in .o files on Mac where the .o files

lldb/tools/lldb-dap/Handler/LaunchRequestHandler.cpp

@@ -30,6 +30,12 @@ Error LaunchRequestHandler::Run(const LaunchRequestArguments &arguments) const {
   dap.SetConfiguration(arguments.configuration, /*is_attach=*/false);
   dap.last_launch_request = arguments;
 
+  // When lldb-dap is reused, the ID of each session is passed in by a request
+  // argument, rather than through the environment variable. Set it into the
+  // environment variable, so that the downstream code can pick it up from there.
+  // See T215634007 for more details.
+  UpdateVSCodeSessionID(arguments, dap);
+
   PrintWelcomeMessage();
 
   // This is a hack for loading DWARF in .o files on Mac where the .o files

lldb/tools/lldb-dap/LLDBUtils.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "LLDBUtils.h"
+#include "DAP.h"
 #include "JSONUtils.h"
 #include "lldb/API/SBCommandInterpreter.h"
 #include "lldb/API/SBCommandReturnObject.h"
@@ -19,6 +20,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/raw_ostream.h"
+#include "Protocol/ProtocolRequests.h"
 
 #include <cstring>
 #include <mutex>
@@ -259,4 +261,20 @@ lldb::SBLineEntry GetLineEntryForAddress(lldb::SBTarget &target,
   return sc.GetLineEntry();
 }
 
+void UpdateVSCodeSessionID(const protocol::LaunchRequestArguments arguments, DAP &dap) {
+  if (!arguments.__sessionId.has_value() || arguments.__sessionId->empty())
+    return; // No session ID provided.
+  dap.configuration.initCommands.insert(
+      dap.configuration.initCommands.begin(),
+      "script os.environ['VSCODE_DEBUG_SESSION_ID'] = '" + *arguments.__sessionId + "'");
+}
+
+void UpdateVSCodeSessionID(const protocol::AttachRequestArguments arguments, DAP &dap) {
+  if (!arguments.__sessionId.has_value() || arguments.__sessionId->empty())
+    return; // No session ID provided.
+  dap.configuration.initCommands.insert(
+      dap.configuration.initCommands.begin(),
+      "script os.environ['VSCODE_DEBUG_SESSION_ID'] = '" + *arguments.__sessionId + "'");
+}
+
 } // namespace lldb_dap

lldb/tools/lldb-dap/LLDBUtils.h

@@ -22,6 +22,7 @@
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "llvm/Support/raw_ostream.h"
+#include "Protocol/ProtocolRequests.h"
 #include <chrono>
 #include <string>
 
@@ -248,6 +249,31 @@ llvm::Error ToError(const lldb::SBError &error);
 /// Provides the string value if this data structure is a string type.
 std::string GetStringValue(const lldb::SBStructuredData &data);
 
+/// Get the vscode session id from the request arguments.
+/// If exists, insert the fetched session id into dap init_commands.
+///
+/// \param[in] arguments
+///     The LaunchRequestArguments 
+///
+/// \param[in] dap
+///     The DAP session
+///
+/// \return
+///     void
+void UpdateVSCodeSessionID(const protocol::LaunchRequestArguments arguments, DAP &dap);
+
+/// Get the vscode session id from the request arguments.
+/// If exists, insert the fetched session id into dap init_commands.
+///
+/// \param[in] arguments
+///     The AttachRequestArguments
+///
+/// \param[in] dap
+///     The DAP session
+///
+/// \return
+///     void
+void UpdateVSCodeSessionID(const protocol::AttachRequestArguments arguments, DAP &dap);
 } // namespace lldb_dap
 
 #endif

lldb/tools/lldb-dap/Protocol/ProtocolRequests.cpp

@@ -277,6 +277,7 @@ bool fromJSON(const json::Value &Params, LaunchRequestArguments &LRA,
          O.mapOptional("shellExpandArguments", LRA.shellExpandArguments) &&
 
          O.mapOptional("runInTerminal", LRA.runInTerminal) &&
+         O.mapOptional("__sessionId", LRA.__sessionId) &&
          parseEnv(Params, LRA.env, P);
 }
 
@@ -289,7 +290,8 @@ bool fromJSON(const json::Value &Params, AttachRequestArguments &ARA,
          O.mapOptional("waitFor", ARA.waitFor) &&
          O.mapOptional("gdb-remote-port", ARA.gdbRemotePort) &&
          O.mapOptional("gdb-remote-hostname", ARA.gdbRemoteHostname) &&
-         O.mapOptional("coreFile", ARA.coreFile);
+         O.mapOptional("coreFile", ARA.coreFile) &&
+         O.mapOptional("__sessionId", ARA.__sessionId);
 }
 
 bool fromJSON(const json::Value &Params, ContinueArguments &CA, json::Path P) {

lldb/tools/lldb-dap/Protocol/ProtocolRequests.h

@@ -302,6 +302,9 @@ struct LaunchRequestArguments {
   /// debugging interactive command line programs.
   bool runInTerminal = false;
 
+  /// VSCode session ID for debugging session tracking.
+  std::optional<std::string> __sessionId;
+
   /// @}
 };
 bool fromJSON(const llvm::json::Value &, LaunchRequestArguments &,
@@ -347,6 +350,9 @@ struct AttachRequestArguments {
   /// Path to the core file to debug.
   std::string coreFile;
 
+  /// VSCode session ID for debugging session tracking.
+  std::optional<std::string> __sessionId;
+
   /// @}
 };
 bool fromJSON(const llvm::json::Value &, AttachRequestArguments &,