Add the ssh-shell task type to allow defining/running tasks that run commands on a remote host

package.json

@@ -190,8 +190,24 @@
                     ]
                 }
             }
+        },
+        "taskDefinitions": [
+            {
+                "type": "ssh-shell",
+                "properties": {
+                    "host": {
+                        "type": "string",
+                        "description": "The configuration (name) to run this task on"
+                    },
+                    "command": {
+                        "type": "string",
+                        "description": "The command to run on the server"
                     }
                 },
+                "required": ["host", "command"]
+            }
+        ]
+    },
     "scripts": {
         "vscode:prepublish": "yarn run build && cd webview && yarn run build",
         "build": "webpack --mode production",

src/extension.ts

@@ -1,7 +1,7 @@
 
 import * as vscode from 'vscode';
 import { loadConfigs } from './config';
-import { FileSystemConfig, invalidConfigName } from './fileSystemConfig';
+import { FileSystemConfig } from './fileSystemConfig';
 import { FileSystemRouter } from './fileSystemRouter';
 import { Logging } from './logging';
 import { Manager } from './manager';
@@ -54,6 +54,7 @@ export function activate(context: vscode.ExtensionContext) {
 
   subscribe(vscode.workspace.registerFileSystemProvider('ssh', new FileSystemRouter(manager), { isCaseSensitive: true }));
   subscribe(vscode.window.createTreeView('sshfs-configs', { treeDataProvider: manager, showCollapseAll: true }));
+  subscribe(vscode.tasks.registerTaskProvider('ssh-shell', manager));
 
   async function pickAndClick(func: (name: string) => void, name?: string, activeOrNot?: boolean) {
     name = name || await pickConfig(manager, activeOrNot);

src/manager.ts

@@ -4,7 +4,7 @@ import * as vscode from 'vscode';
 import { configMatches, getConfig, getConfigs, loadConfigs, loadConfigsRaw, UPDATE_LISTENERS } from './config';
 import { FileSystemConfig, getGroups } from './fileSystemConfig';
 import { Logging } from './logging';
-import { SSHPseudoTerminal } from './pseudoTerminal';
+import { createTaskTerminal, SSHPseudoTerminal } from './pseudoTerminal';
 import { catchingPromise, toPromise } from './toPromise';
 import { Navigation } from './webviewMessages';
 
@@ -58,7 +58,7 @@ interface Connection {
   pendingUserCount: number;
 }
 
-export class Manager implements vscode.TreeDataProvider<string | FileSystemConfig> {
+export class Manager implements vscode.TreeDataProvider<string | FileSystemConfig>, vscode.TaskProvider {
   public onDidChangeTreeData: vscode.Event<string | null>;
   protected connections: Connection[] = [];
   protected pendingConnections: { [name: string]: Promise<Connection> } = {};
@@ -281,6 +281,31 @@ export class Manager implements vscode.TreeDataProvider<string | FileSystemConfi
     }
     return [...matching, ...groups.sort()];
   }
+  /* TaskProvider */
+  public provideTasks(token?: vscode.CancellationToken | undefined): vscode.ProviderResult<vscode.Task[]> {
+    return [];
+  }
+  public async resolveTask(task: vscode.Task, token?: vscode.CancellationToken | undefined): Promise<vscode.Task> {
+    const { host, command } = task.definition as { host?: string, command?: string };
+    if (!host) throw new Error('Missing field \'host\' for ssh-shell task');
+    if (!command) throw new Error('Missing field \'command\' for ssh-shell task');
+    const config = getConfig(host);
+    if (!config) throw new Error(`No configuration with the name '${host}' found for ssh-shell task`);
+    return new vscode.Task(
+      task.definition,
+      vscode.TaskScope.Workspace,
+      `SSH Task for ${host}`,
+      'ssh',
+      new vscode.CustomExecution(async () => {
+        const connection = await this.createConnection(host);
+        return createTaskTerminal({
+          command,
+          client: connection.client,
+          config: connection.actualConfig,
+        })
+      })
+    )
+  }
   /* Commands (stuff for e.g. context menu for ssh-configs tree) */
   public commandDisconnect(target: string | FileSystemConfig) {
     if (typeof target === 'object') target = target.name;

src/pseudoTerminal.ts

@@ -5,7 +5,7 @@ import { FileSystemConfig } from "./fileSystemConfig";
 import { toPromise } from "./toPromise";
 
 const [HEIGHT, WIDTH] = [480, 640];
-const PSEUDO_TY_OPTIONS: PseudoTtyOptions = {
+const PSEUDO_TTY_OPTIONS: PseudoTtyOptions = {
     height: HEIGHT, width: WIDTH,
 };
 
@@ -13,24 +13,29 @@ export interface SSHPseudoTerminal extends vscode.Pseudoterminal {
     onDidClose: vscode.Event<number>; // Redeclaring that it isn't undefined
     config: FileSystemConfig;
     client: Client;
-    channel: ClientChannel;
+    /** Could be undefined if it only gets created during psy.open() instead of beforehand */
+    channel?: ClientChannel;
 }
 
 export async function createTerminal(client: Client, config: FileSystemConfig): Promise<SSHPseudoTerminal> {
-    const channel = await toPromise<ClientChannel | undefined>(cb => client.shell(PSEUDO_TY_OPTIONS, cb));
+    const channel = await toPromise<ClientChannel | undefined>(cb => client.shell(PSEUDO_TTY_OPTIONS, cb));
     if (!channel) throw new Error('Could not create remote terminal');
     const onDidWrite = new vscode.EventEmitter<string>();
     onDidWrite.fire(`Connecting to ${config.label || config.name}...\n`);
     (channel as Readable).on('data', chunk => onDidWrite.fire(chunk.toString()));
+    channel.stderr.on('data', chunk => onDidWrite.fire(chunk.toString()));
     const onDidClose = new vscode.EventEmitter<number>();
     channel.on('exit', onDidClose.fire);
     // Hopefully the exit event fires first
-    channel.on('close', () => onDidClose.fire(1));
+    channel.on('close', () => onDidClose.fire(0));
     const pseudo: SSHPseudoTerminal = {
         config, client, channel,
         onDidWrite: onDidWrite.event,
         onDidClose: onDidClose.event,
         close() {
+            channel.signal('INT');
+            channel.signal('SIGINT');
+            channel.write('\x03');
             channel.close();
         },
         open(dims) {
@@ -46,3 +51,56 @@ export async function createTerminal(client: Client, config: FileSystemConfig):
     };
     return pseudo;
 }
+
+export interface TaskTerminalOptions {
+    client: Client;
+    config: FileSystemConfig;
+    command: string;
+}
+
+export async function createTaskTerminal(options: TaskTerminalOptions): Promise<SSHPseudoTerminal> {
+    const { client, config, command } = options;
+    const onDidWrite = new vscode.EventEmitter<string>();
+    onDidWrite.fire(`Connecting to ${config.label || config.name}...\n`);
+    const onDidClose = new vscode.EventEmitter<number>();
+    let channel: ClientChannel;
+    const pseudo: SSHPseudoTerminal = {
+        config, client,
+        onDidWrite: onDidWrite.event,
+        onDidClose: onDidClose.event,
+        close() {
+            channel?.signal('INT');
+            channel?.signal('SIGINT');
+            channel?.write('\x03');
+            channel?.close();
+        },
+        open(dims) {
+            onDidWrite.fire(`Running command: ${command}\n`);
+            (async () => {
+                const ch = await toPromise<ClientChannel | undefined>(cb => client.exec(command, {
+                    pty: { ...PSEUDO_TTY_OPTIONS, cols: dims?.columns, rows: dims?.rows }
+                }, cb));
+                if (!ch) {
+                    onDidWrite.fire(`Could not create SSH channel, running task failed\n`);
+                    onDidClose.fire(1);
+                    return;
+                }
+                pseudo.channel = channel = ch;
+                channel.on('exit', onDidClose.fire);
+                channel.on('close', () => onDidClose.fire(0));
+                (channel as Readable).on('data', chunk => onDidWrite.fire(chunk.toString()));
+                channel.stderr.on('data', chunk => onDidWrite.fire(chunk.toString()));
+            })().catch(e => {
+                onDidWrite.fire(`Error starting process over SSH:\n${e}\n`);
+                onDidClose.fire(1);
+            });
+        },
+        setDimensions(dims) {
+            channel?.setWindow(dims.rows, dims.columns, HEIGHT, WIDTH);
+        },
+        handleInput(data) {
+            channel?.write(data);
+        },
+    };
+    return pseudo;
+}