Environment support for multiple hosts or ports - Redo

src/EnvironmentManager.ts

@@ -0,0 +1,192 @@
+import { readFile, writeFile, mkdir, access, unlink, readdir } from "fs/promises";
+import { join } from "path";
+import { lmsConfigFolder } from "./lmstudioPaths.js";
+import { z } from "zod";
+import { type SimpleLogger } from "@lmstudio/lms-common";
+import { type CreateClientArgs } from "./createClient.js";
+
+const environmentConfigSchema = z.object({
+  name: z.string(),
+  host: z.string(),
+  port: z.number().int().min(0).max(65535),
+  description: z.string().optional(),
+});
+export type EnvironmentConfig = z.infer<typeof environmentConfigSchema>;
+
+export const DEFAULT_LOCAL_ENVIRONMENT_NAME = "local";
+
+const DEFAULT_ENVIRONMENT_CONFIG: EnvironmentConfig = {
+  name: DEFAULT_LOCAL_ENVIRONMENT_NAME,
+  host: "localhost",
+  // TODO: This will have to change based on which port the current user is running their local
+  // server on and consider multiple users on the same machine using different ports
+  port: 1234,
+  description: "Default local environment",
+};
+
+export class EnvironmentManager {
+  private environmentsDir: string;
+  private currentEnvFile: string;
+
+  public constructor(
+    private readonly logger: SimpleLogger,
+    private readonly createClientArgs: CreateClientArgs | undefined = undefined,
+  ) {
+    const configDir = lmsConfigFolder;
+    this.environmentsDir = join(configDir, "environments");
+    this.currentEnvFile = join(configDir, "current-env");
+  }
+
+  private async ensureDirExists(): Promise<void> {
+    await mkdir(this.environmentsDir, { recursive: true });
+  }
+
+  /**
+   * Adds a new environment configuration.
+   */
+  public async addEnvironment(config: EnvironmentConfig): Promise<void> {
+    await this.ensureDirExists();
+    const envPath = join(this.environmentsDir, `${config.name}.json`);
+    try {
+      await access(envPath);
+      throw new Error(`Environment ${config.name} already exists.`);
+    } catch {
+      await writeFile(envPath, JSON.stringify(config, null, 2), "utf-8");
+    }
+  }
+
+  /**
+   * Removes an existing environment configuration.
+   */
+  public async removeEnvironment(name: string): Promise<void> {
+    const envPath = join(this.environmentsDir, `${name}.json`);
+    try {
+      await unlink(envPath);
+      // Check if this was the current environment
+      try {
+        const currentEnv = await readFile(this.currentEnvFile, "utf-8");
+        if (currentEnv === name) {
+          await writeFile(this.currentEnvFile, DEFAULT_LOCAL_ENVIRONMENT_NAME, "utf-8");
+        }
+      } catch (error) {
+        if (error instanceof Error && "code" in error && error.code === "ENOENT") {
+          await writeFile(this.currentEnvFile, DEFAULT_LOCAL_ENVIRONMENT_NAME, "utf-8");
+        } else {
+          // Re-throw other types of errors
+          throw error;
+        }
+      }
+    } catch (error) {
+      if (error instanceof Error && "code" in error && error.code === "ENOENT") {
+        throw new Error(`Environment ${name} does not exist.`);
+      } else {
+        throw new Error(`Failed to remove environment ${name}: ${(error as Error).message}`);
+      }
+    }
+  }
+
+  /**
+   * Sets the current environment by name.
+   */
+  public async setCurrentEnvironment(name: string): Promise<void> {
+    if (name === "local") {
+      // Special case for local environment
+      await writeFile(this.currentEnvFile, DEFAULT_LOCAL_ENVIRONMENT_NAME, "utf-8");
+      return;
+    }
+    const envPath = join(this.environmentsDir, `${name}.json`);
+    try {
+      const data = await readFile(envPath, "utf-8");
+      environmentConfigSchema.parse(JSON.parse(data)); // Validate schema
+      await writeFile(this.currentEnvFile, name, "utf-8");
+    } catch {
+      throw new Error(`Environment ${name} does not exist.`);
+    }
+  }
+
+  /**
+   * Gets the current environment configuration.
+   */
+  public async getCurrentEnvironment(): Promise<EnvironmentConfig> {
+    if (this.createClientArgs !== undefined) {
+      if (this.createClientArgs.host !== undefined || this.createClientArgs.port !== undefined) {
+        // Check if --host and --port are provided in createClientOpts
+        // If so, return a temporary environment config
+        const tempEnv: EnvironmentConfig = {
+          name: "temporary",
+          host: this.createClientArgs.host ?? "localhost",
+          port: this.createClientArgs.port ?? 1234,
+        };
+        return tempEnv;
+      }
+    }
+
+    let envName: string;
+
+    // Check if LMS_ENV is set in the environment variables
+    // This takes precedence over the currentEnvFile
+    if (
+      process.env.LMS_ENV !== undefined &&
+      process.env.LMS_ENV !== "undefined" &&
+      process.env.LMS_ENV !== "null"
+    ) {
+      envName = process.env.LMS_ENV;
+    } else {
+      try {
+        envName = (await readFile(this.currentEnvFile, "utf-8")).trim();
+      } catch {
+        envName = DEFAULT_LOCAL_ENVIRONMENT_NAME;
+      }
+    }
+    if (envName === undefined || envName === "" || envName === DEFAULT_LOCAL_ENVIRONMENT_NAME) {
+      return DEFAULT_ENVIRONMENT_CONFIG;
+    }
+
+    const env = await this.tryGetEnvironment(envName);
+    if (env === undefined) {
+      this.logger.warn(`Environment ${envName} not found, falling back to local.`);
+      await writeFile(this.currentEnvFile, DEFAULT_LOCAL_ENVIRONMENT_NAME, "utf-8");
+      return DEFAULT_ENVIRONMENT_CONFIG;
+    }
+
+    return env;
+  }
+
+  /**
+   * Lists all available environment configurations.
+   */
+  public async getAllEnvironments(): Promise<EnvironmentConfig[]> {
+    await this.ensureDirExists();
+    const files = await readdir(this.environmentsDir);
+    const environments: EnvironmentConfig[] = [DEFAULT_ENVIRONMENT_CONFIG];
+    for (const file of files) {
+      if (file.endsWith(".json")) {
+        try {
+          const data = await readFile(join(this.environmentsDir, file), "utf-8");
+          const parsed = environmentConfigSchema.parse(JSON.parse(data));
+          environments.push(parsed);
+        } catch (error) {
+          this.logger.warn(`Failed to load environment from ${file}: ${(error as Error).message}`);
+        }
+      }
+    }
+    return environments;
+  }
+
+  /**
+   * Tries to get an environment configuration by name. Returns undefined if not found.
+   */
+  public async tryGetEnvironment(name: string): Promise<EnvironmentConfig | undefined> {
+    if (name === DEFAULT_LOCAL_ENVIRONMENT_NAME) {
+      return DEFAULT_ENVIRONMENT_CONFIG; // Return default local environment
+    }
+    await this.ensureDirExists();
+    const envPath = join(this.environmentsDir, `${name}.json`);
+    try {
+      const data = await readFile(envPath, "utf-8");
+      return environmentConfigSchema.parse(JSON.parse(data));
+    } catch {
+      return undefined; // Environment does not exist
+    }
+  }
+}

src/createClient.ts

@@ -1,14 +1,15 @@
 import { Option, type Command, type OptionValues } from "@commander-js/extra-typings";
 import { apiServerPorts, text, type SimpleLogger } from "@lmstudio/lms-common";
 import { LMStudioClient, type LMStudioClientConstructorOpts } from "@lmstudio/sdk";
-import chalk from "chalk";
 import { spawn } from "child_process";
 import { randomBytes } from "crypto";
 import { readFile } from "fs/promises";
 import { exists } from "./exists.js";
 import { appInstallLocationFilePath, lmsKey2Path } from "./lmstudioPaths.js";
 import { type LogLevelArgs } from "./logLevel.js";
 import { createRefinedNumberParser } from "./types/refinedNumber.js";
+import { DEFAULT_LOCAL_ENVIRONMENT_NAME, EnvironmentManager } from "./EnvironmentManager.js";
+import chalk from "chalk";
 /**
  * Checks if the HTTP server is running.
  */
@@ -78,13 +79,22 @@ export function addCreateClientOptions<
           (default), the last used port will be used; otherwise, 1234 will be used.
         `,
       ).argParser(createRefinedNumberParser({ integer: true, min: 0, max: 65535 })),
+    )
+    .addOption(
+      new Option(
+        "--env <env>",
+        text`
+          If you wish to connect to a remote LM Studio instance, specify the env here
+          `,
+      ),
     );
 }
 
-interface CreateClientArgs {
+export interface CreateClientArgs {
   yes?: boolean;
   host?: string;
   port?: number;
+  env?: string;
 }
 
 async function isLocalServerAtPortLMStudioServerOrThrow(port: number) {
@@ -151,12 +161,48 @@ export async function createClient(
   args: CreateClientArgs & LogLevelArgs,
   _opts: CreateClientOpts = {},
 ) {
-  let { host, port } = args;
-  let isRemote = true;
-  if (host === undefined) {
-    isRemote = false;
-    host = "127.0.0.1";
-  } else if (host.includes("://")) {
+  const envManager = new EnvironmentManager(logger, args);
+  let currentEnv = await envManager.getCurrentEnvironment();
+  let host = currentEnv.host;
+  let port = currentEnv.port;
+
+  // If the user has provided both host and port and env, we will throw an error.
+  if (args.host !== undefined && args.port !== undefined && args.env !== undefined) {
+    logger.error(
+      text`
+        You cannot specify both --host/--port and --env at the same time. Please choose one method
+        to specify the server to connect to.
+        `,
+    );
+    process.exit(1);
+  }
+  if (args.env !== undefined) {
+    // If the user specified an environment, we will override the current environment with the specified one.
+    const specificedEnv = await envManager.tryGetEnvironment(args.env);
+    if (specificedEnv === undefined) {
+      logger.error(`Environment '${args.env}' not found`);
+      process.exit(1);
+    }
+    host = specificedEnv.host;
+    port = specificedEnv.port;
+    currentEnv = specificedEnv;
+  }
+  let isRemote = currentEnv.name !== DEFAULT_LOCAL_ENVIRONMENT_NAME;
+
+  // If host and port are provided, they take precedence over the environment.
+  if (args.host !== undefined) {
+    host = args.host;
+
+    // If host is anything but remote, set isRemote to false
+    if (host === "127.0.0.1" || host === "localhost" || host === "0.0.0.0") {
+      isRemote = false;
+    }
+  }
+  if (args.port !== undefined) {
+    port = args.port;
+  }
+
+  if (host.includes("://")) {
     logger.error("Host should not include the protocol.");
     process.exit(1);
   } else if (host.includes(":")) {
@@ -202,7 +248,7 @@ export async function createClient(
       }
     }
   }
-  if (port === undefined && host === "127.0.0.1") {
+  if (isRemote === false) {
     // We will now attempt to connect to the local API server.
     const localPort = await tryFindLocalAPIServer();
 
@@ -242,11 +288,6 @@ export async function createClient(
 
     logger.error("");
   }
-
-  if (port === undefined) {
-    port = 1234;
-  }
-
   logger.debug(`Connecting to server at ${host}:${port}`);
   if (!(await checkHttpServer(logger, port, host))) {
     logger.error(

src/index.ts

@@ -16,6 +16,7 @@ import { server } from "./subcommands/server.js";
 import { status } from "./subcommands/status.js";
 import { unload } from "./subcommands/unload.js";
 import { printVersion, version } from "./subcommands/version.js";
+import { env } from "./subcommands/env.js";
 
 if (process.argv.length === 2) {
   printVersion();
@@ -46,6 +47,7 @@ program.addCommand(push);
 
 program.commandsGroup("System Management:");
 program.addCommand(bootstrap);
+progrma.addCommand(env);
 program.addCommand(flags);
 program.addCommand(log);
 program.addCommand(status);

src/lmstudioPaths.ts

@@ -2,6 +2,7 @@ import { findLMStudioHome } from "@lmstudio/lms-common-server";
 import { join } from "path";
 
 const lmstudioHome = findLMStudioHome();
+export const lmsConfigFolder = join(lmstudioHome, "lms"); //TODO: Temporary path
 export const pluginsFolderPath = join(lmstudioHome, "extensions", "plugins");
 export const lmsKey2Path = join(lmstudioHome, ".internal", "lms-key-2");
 export const cliPrefPath = join(lmstudioHome, ".internal", "cli-pref.json");