From ba0cc0fa2cd639ed7b73a7d1703add5b5c98039c Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Fri, 22 Aug 2025 21:14:33 +0530
Subject: [PATCH] save_as_session: Implement saving of foreground process for
 remote shells run via ssh kitten

---
 kittens/ssh/utils.py | 23 +++++++++++++++++++++--
 kitty/window.py      |  6 +++++-
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/kittens/ssh/utils.py b/kittens/ssh/utils.py
index b3b7649a8..f7d3ad34d 100644
--- a/kittens/ssh/utils.py
+++ b/kittens/ssh/utils.py
@@ -32,7 +32,7 @@ def ssh_options() -> dict[str, str]:
     ans: dict[str, str] = {}
     pos = 0
     while True:
-        pos = raw.find('[', pos)
+        pos = raw.find('[', pos)  # ]
         if pos < 0:
             break
         num = 1
@@ -41,7 +41,7 @@ def ssh_options() -> dict[str, str]:
             epos += 1
             if raw[epos] not in '[]':
                 continue
-            num += 1 if raw[epos] == '[' else -1
+            num += 1 if raw[epos] == '[' else -1  # ]
         q = raw[pos+1:epos]
         pos = epos
         if len(q) < 2 or q[0] != '-':
@@ -81,6 +81,25 @@ def patch_cmdline(key: str, val: str, argv: list[str]) -> None:
     argv.insert(idx + 1, f'--kitten={key}={val}')
 
 
+def remove_env_var_from_cmdline(key: str, argv: list[str]) -> None:
+    while True:
+        for i, arg in enumerate(tuple(argv)):
+            if arg.startswith(f'--kitten=env={key}='):
+                del argv[i]
+                break
+            elif i > 0 and argv[i-1] == '--kitten' and (arg.startswith(f'env={key}=') or arg.startswith(f'env {key}=')):
+                del argv[i-1:i+1]
+                break
+        else:
+            break
+
+
+def set_single_env_var_in_cmdline(key: str, val: str, argv: list[str]) -> None:
+    remove_env_var_from_cmdline(key, argv)
+    idx = argv.index('ssh')
+    argv.insert(idx+1, f'--kitten=env={key}={val}')
+
+
 def set_cwd_in_cmdline(cwd: str, argv: list[str]) -> None:
     patch_cmdline('cwd', cwd, argv)
 
diff --git a/kitty/window.py b/kitty/window.py
index 0284b643b..35b8c3f0d 100644
--- a/kitty/window.py
+++ b/kitty/window.py
@@ -2028,10 +2028,14 @@ class Window:
                         cmd[0] = abspath_of_exe(pid)
             kssh_cmdline = self.ssh_kitten_cmdline()
             if kssh_cmdline:
-                from kittens.ssh.utils import set_cwd_in_cmdline
+                from kittens.ssh.utils import remove_env_var_from_cmdline, set_cwd_in_cmdline, set_single_env_var_in_cmdline
+                remove_env_var_from_cmdline('KITTY_SI_RUN_COMMAND_AT_STARTUP', kssh_cmdline)
                 if self.at_prompt:
                     if self.screen.last_reported_cwd:
                         set_cwd_in_cmdline(path_from_osc7_url(self.screen.last_reported_cwd), kssh_cmdline)
+                else:
+                    if self.last_cmd_cmdline:
+                        set_single_env_var_in_cmdline('KITTY_SI_RUN_COMMAND_AT_STARTUP', self.last_cmd_cmdline, kssh_cmdline)
                 unserialize_data['cmd_at_shell_startup'] = kssh_cmdline
             elif not self.at_prompt:
                 if self.last_cmd_cmdline: