[commands] provide command args directly (#107)

* [commands] provide command args directly

* fix lint

* handle null case

* fixed bugs

modules/extensions/src/api/experimental/commands.ts

@@ -1,5 +1,14 @@
-import { extensionPort } from "../../util/comlink";
-import { CreateCommand, CommandProxy, ContributionType } from "../../commands";
+import { extensionPort, proxy } from "../../util/comlink";
+import {
+  CreateCommand,
+  CommandProxy,
+  ContributionType,
+  Command,
+  CommandFnArgs,
+  CommandSymbol,
+  CommandArgs,
+  isCommandProxy,
+} from "../../commands";
 
 export interface AddCommandArgs {
   /**
@@ -27,16 +36,28 @@ export interface AddCommandArgs {
  */
 export function add({ id, contributions, command }: AddCommandArgs) {
   if (typeof command === "function") {
+    let createCommand = proxy(async (cmdFnArgs: CommandFnArgs) => {
+      const cmd = await command(cmdFnArgs);
+
+      if (!cmd) {
+        return null;
+      }
+
+      return isCommandProxy(cmd) ? cmd : Command(cmd);
+    });
+
     extensionPort.experimental.commands.registerCreateCommand(
       { commandId: id, contributions },
-      command
+      createCommand
     );
   } else {
+    let createCommand = proxy(async () => {
+      return isCommandProxy(command) ? command : Command(command);
+    });
+
     extensionPort.experimental.commands.registerCreateCommand(
       { commandId: id, contributions },
-      async () => ({
-        ...command,
-      })
+      createCommand
     );
   }
 }
@@ -57,6 +78,14 @@ export function registerCreate(
 ): void {
   extensionPort.experimental.commands.registerCreateCommand(
     data,
-    createCommand
+    async (args: CommandFnArgs) => {
+      const cmd = await createCommand(args);
+
+      if (!cmd) {
+        return null;
+      }
+
+      return isCommandProxy(cmd) ? cmd : Command(cmd);
+    }
   );
 }

modules/extensions/src/commands/index.ts

@@ -1,3 +1,4 @@
+import { ProxyMarked } from "comlink";
 import { proxy } from "../util/comlink";
 
 /**
@@ -31,11 +32,18 @@ export type CommandFnArgs = {
   path: SerializableValue[];
 };
 
-export type CommandsFn = (args: CommandFnArgs) => Promise<Array<CommandProxy>>;
+export type CommandsFn = (
+  args: CommandFnArgs
+) => Promise<Array<CommandProxy | CommandArgs>>;
 
 export type CreateCommand = (
   args: CommandFnArgs
-) => CommandProxy | Promise<CommandProxy>;
+) =>
+  | CommandProxy
+  | Promise<CommandProxy>
+  | CommandArgs
+  | Promise<CommandArgs>
+  | null;
 
 export type Run = () => any;
 
@@ -85,9 +93,9 @@ export type ContextCommandArgs = BaseCommandArgs & {
 export type CommandArgs = ActionCommandArgs | ContextCommandArgs;
 
 /**
- * This validates a command. We make sure that exactly one of `commands` or `run` is defined, and every other argument is serializable.
+ * This validates a CommandArgs object. We make sure that exactly one of `commands` or `run` is defined, and every other argument is serializable.
  */
-function validateCommand(cmdArgs: unknown): asserts cmdArgs is CommandArgs {
+function validateCommandArgs(cmdArgs: unknown): asserts cmdArgs is CommandArgs {
   // Make sure cmdArgs is object
   if (typeof cmdArgs !== "object") {
     throw new Error("Command arguments must be an object");
@@ -134,31 +142,101 @@ function validateCommand(cmdArgs: unknown): asserts cmdArgs is CommandArgs {
   }
 }
 
-export function Command(cmdArgs: CommandArgs) {
-  validateCommand(cmdArgs);
+export const CommandSymbol = Symbol("Command");
+
+export function isCommandProxy(cmd: object): cmd is CommandProxy {
+  return CommandSymbol in cmd && cmd[CommandSymbol] === true;
+}
+
+/**
+ * This function validates a command and wraps it in a proxy so that it can be sent over the wire
+ *
+ * It:
+ * - Validates the command's arguments, separates serializable and non-serializable arguments
+ * - Wraps the command in a proxy so that it can be sent over the wire
+ * - Wraps the command's `commands` function to ensure that all subcommands are also Command() wrapped
+ * - Adds a symbol to the command to identify a wrapped command
+ */
+export function Command(cmdArgs: CommandArgs): CommandProxy {
+  // If the command is already wrapped, just return it.
+  // This is to prevent accidental double-wrapping
+  if (isCommandProxy(cmdArgs)) {
+    throw new Error("Command is already wrapped");
+  }
+
+  // Validate the command's arguments
+  validateCommandArgs(cmdArgs);
 
   if ("commands" in cmdArgs) {
     const { commands, ...props } = cmdArgs;
-    return proxy({
+
+    let cmdProxy: CommandProxy = proxy({
       data: {
         ...props,
         type: "context",
       },
       commands: async (args: CommandFnArgs) => {
-        return proxy(await commands(args));
+        // Compute subcommands
+        let subCmds = await commands(args);
+
+        // While we expect commands() to return an array, we don't want to throw an error if it doesn't.
+        if (!subCmds || !Array.isArray(subCmds)) {
+          return proxy([]);
+        }
+
+        const commandProxyArray: Array<CommandProxy> = subCmds.map((subCmd) => {
+          // Subcommands can be either a CommandArgs or a CommandProxy.
+          // If it's already a wrapped command, just return it.
+          if (isCommandProxy(subCmd)) {
+            return subCmd;
+          }
+
+          // Otherwise, wrap it in Command()
+          return Command(subCmd);
+        });
+
+        // Return the subcommands as a proxy array
+        return proxy(commandProxyArray);
       },
+
+      // Attach CommandSymbol to identify a wrapped command
+      [CommandSymbol]: true,
     });
+
+    return cmdProxy;
   } else {
     const { run, ...props } = cmdArgs;
 
-    return proxy({
+    let cmdProxy: CommandProxy = proxy({
       data: {
         ...props,
         type: "action",
       },
       run,
+      // Attach CommandSymbol to identify a wrapped command
+      [CommandSymbol]: true,
     });
+
+    return cmdProxy;
   }
 }
 
-export type CommandProxy = ReturnType<typeof Command>;
+export type CommandProxy =
+  | ({
+      data: {
+        type: "action";
+        label: string;
+        description?: string;
+        icon?: string;
+      };
+      run?: Run;
+    } & ProxyMarked & { [CommandSymbol]: true })
+  | ({
+      data: {
+        type: "context";
+        label: string;
+        description?: string;
+        icon?: string;
+      };
+      commands?: (args: CommandFnArgs) => Promise<Array<CommandProxy>>;
+    } & ProxyMarked & { [CommandSymbol]: true });

modules/extensions/src/types/index.ts

@@ -217,7 +217,7 @@ export type ExperimentalAPI = {
         commandId: string;
         contributions: Array<string>;
       },
-      create: CreateCommand
+      create: (createArgs: CommandFnArgs) => Promise<CommandProxy | null>
     ) => void;
   };
 };