Have outleap-agent forward sys.argv[1:]

src/outleap/client.py

@@ -6,6 +6,7 @@
 import enum
 import logging
 import os
+import sys
 import uuid
 from types import TracebackType
 from typing import *
@@ -34,6 +35,10 @@ def __init__(self, protocol: AbstractLEAPProtocol):
         self.cmd_pump: Optional[str] = None
         # Process ID of the viewer connecting to us
         self.viewer_pid: Optional[int] = None
+        # Arguments the script the client belongs to was launched with
+        # This is necessary to distinguish between LEAP scripts launched with
+        # the --leap arguments and those launched as puppetry plugins.
+        self.launch_args: Optional[List[str]] = None
         # Map of req id -> future held by requester to send responses to
         self._reply_futs: Dict[uuid.UUID, asyncio.Future] = {}
         self._pump_listeners: Dict[str, ListenerDetails] = collections.defaultdict(ListenerDetails)
@@ -72,10 +77,13 @@ async def connect(self) -> None:
             self._reply_pump = welcome_message["pump"]
             self.cmd_pump = welcome_message["data"]["command"]
             self.viewer_pid = welcome_message["data"].get("process_id")
+            self.launch_args = welcome_message["data"].get("args")
             if self.viewer_pid is None:
                 # handshake didn't include the process ID, assume our own parent
                 # process ID is the viewer's process ID.
                 self.viewer_pid = os.getppid()
+            if self.launch_args is None:
+                self.launch_args = sys.argv[1:]
 
             self._connection_status = ConnectionStatus.CONNECTED
             self._start_message_pump()

src/outleap/scripts/agent.py

@@ -10,6 +10,7 @@
 import asyncio
 import multiprocessing
 import os
+import sys
 
 from outleap import LEAPProtocol
 from outleap.utils import connect_stdin_stdout
@@ -26,9 +27,11 @@ async def amain():
     stdio_reader, stdio_writer = await connect_stdin_stdout()
 
     # Enrich the LEAP handshake with our parent's process ID so
-    # the connection can be tied to a specific client
+    # the connection can be tied to a specific client.
+    # Pass along any arguments we were launched with as well.
     handshake = await LEAPProtocol(stdio_reader, stdio_writer).read_message()
     handshake["data"]["process_id"] = os.getppid()
+    handshake["data"]["args"] = sys.argv[1:]
 
     # Forward along the modified handshake
     serv_protocol = LEAPProtocol(serv_reader, serv_writer)

tests/__init__.py

@@ -51,6 +51,7 @@ def _write_welcome(self):
                 "data": {
                     "command": "cmd_pump",
                     "process_id": 1,
+                    "args": ["foo", "bar"],
                 },
             }
         )

tests/test_client.py

@@ -1,5 +1,6 @@
 import asyncio
 import os
+import sys
 
 from . import BaseClientTest
 
@@ -14,6 +15,7 @@ async def test_connect_process_id(self):
         self._write_welcome()
         await self.client.connect()
         self.assertEqual(1, self.client.viewer_pid)
+        self.assertListEqual(["foo", "bar"], self.client.launch_args)
 
     async def test_connect_no_process_id(self):
         self.protocol.inbound_messages.put_nowait(
@@ -24,6 +26,7 @@ async def test_connect_no_process_id(self):
         )
         await self.client.connect()
         self.assertEqual(os.getppid(), self.client.viewer_pid)
+        self.assertEqual(sys.argv[1:], self.client.launch_args)
 
     async def test_disconnect(self):
         self._write_welcome()