From 7d94bb8c7a4763cad8ba602c6a038f1b3ad3e852 Mon Sep 17 00:00:00 2001
From: Kelvin Schoofs <schoofs.kelvin@gmail.com>
Date: Fri, 19 Mar 2021 17:58:29 +0100
Subject: [PATCH] Add taskCommand

---
 src/fileSystemConfig.ts |  2 ++
 src/manager.ts          | 13 +++++++++++--
 src/pseudoTerminal.ts   |  4 ++--
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/fileSystemConfig.ts b/src/fileSystemConfig.ts
index 83a81cb..7ade3a9 100644
--- a/src/fileSystemConfig.ts
+++ b/src/fileSystemConfig.ts
@@ -98,6 +98,8 @@ export interface FileSystemConfig extends ConnectConfig {
   sftpSudo?: string | boolean;
   /** The command(s) to run when a new SSH terminal gets created. Defaults to `$SHELL`. Internally the command `cd ...` is run first */
   terminalCommand?: string | string[];
+  /** The command(s) to run when a `ssh-shell` gets run. Defaults to the placeholder `$COMMAND`. Internally the command `cd ...` is run first */
+  taskCommand?: string | string[];
   /** The filemode to assign to created files */
   newFileMode?: number | string;
   /** Whether this config was created from an instant connection string. Enables fuzzy matching for e.g. PuTTY, config-by-host, ... */
diff --git a/src/manager.ts b/src/manager.ts
index 55c8b2c..01c04ab 100644
--- a/src/manager.ts
+++ b/src/manager.ts
@@ -283,13 +283,22 @@ export class Manager implements vscode.TaskProvider, vscode.TerminalLinkProvider
       `SSH Task '${task.name}'`,
       'ssh',
       new vscode.CustomExecution(async (resolved: SSHShellTaskOptions) => {
-        const { createTerminal, createTextTerminal } = await import('./pseudoTerminal');
+        const { createTerminal, createTextTerminal, joinCommands } = await import('./pseudoTerminal');
         try {
           if (!resolved.host) throw new Error('Missing field \'host\' in task description');
           if (!resolved.command) throw new Error('Missing field \'command\' in task description');
           const connection = await this.connectionManager.createConnection(resolved.host);
           resolved = await this.replaceTaskVariablesRecursive(resolved, value => this.replaceTaskVariables(value, connection.actualConfig));
-          const { command, workingDirectory } = resolved;
+          let { command, workingDirectory } = resolved;
+          let { taskCommand = '$COMMAND' } = connection.actualConfig;
+          taskCommand = joinCommands(taskCommand)!;
+          if (taskCommand.includes('$COMMAND')) {
+            command = taskCommand.replace(/\$COMMAND/g, command);
+          } else {
+            const message = `The taskCommand '${taskCommand}' is missing the '$COMMAND' placeholder!`;
+            Logging.warning(message, LOGGING_NO_STACKTRACE);
+            command = `echo "Missing '$COMMAND' placeholder"`;
+          }
           //if (workingDirectory) workingDirectory = this.getRemotePath(config, workingDirectory);
           this.connectionManager.update(connection, con => con.pendingUserCount++);
           const pty = await createTerminal({
diff --git a/src/pseudoTerminal.ts b/src/pseudoTerminal.ts
index 4e34d0f..b774ef5 100644
--- a/src/pseudoTerminal.ts
+++ b/src/pseudoTerminal.ts
@@ -36,10 +36,10 @@ export interface TerminalOptions {
     command?: string;
 }
 
-function joinCommands(commands?: string | string[]): string | undefined {
+export function joinCommands(commands?: string | string[]): string | undefined {
     if (!commands) return undefined;
     if (typeof commands === 'string') return commands;
-    return commands.join(';');
+    return commands.join('; ');
 }
 
 export async function createTerminal(options: TerminalOptions): Promise<SSHPseudoTerminal> {