Skip to content

Commit a00ed1f

Browse files
committed
🤖 refactor: DRY cleanup for Docker runtime implementation
Extract shared utilities and consolidate duplicate code: 1. streamUtils.ts - Extract streamToString helper used by SSH and Docker 2. streamUtils.ts - Extract shescape shell escaping helper used by SSH and Docker 3. initHook.ts - Add runInitHookOnRuntime() shared by SSH and Docker runtimes 4. chatCommands.ts - Use parseRuntimeModeAndHost from common/types/runtime Net reduction: ~79 lines of duplicate code removed while improving maintainability. _Generated with mux_
1 parent 7d3e54f commit a00ed1f

File tree

5 files changed

+166
-245
lines changed

5 files changed

+166
-245
lines changed

src/browser/utils/chatCommands.ts

Lines changed: 37 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ import type {
1616
} from "@/common/types/message";
1717
import type { FrontendWorkspaceMetadata } from "@/common/types/workspace";
1818
import type { RuntimeConfig } from "@/common/types/runtime";
19-
import { RUNTIME_MODE, SSH_RUNTIME_PREFIX, DOCKER_RUNTIME_PREFIX } from "@/common/types/runtime";
19+
import { RUNTIME_MODE, parseRuntimeModeAndHost } from "@/common/types/runtime";
2020
import { CUSTOM_EVENTS, createCustomEvent } from "@/common/constants/events";
2121
import { WORKSPACE_ONLY_COMMANDS } from "@/constants/slashCommands";
2222
import type { Toast } from "@/browser/components/ChatInputToast";
@@ -436,7 +436,9 @@ async function handleForkCommand(
436436
}
437437

438438
/**
439-
* Parse runtime string from -r flag into RuntimeConfig for backend
439+
* Parse runtime string from -r flag into RuntimeConfig for backend.
440+
* Uses shared parseRuntimeModeAndHost for parsing, then converts to RuntimeConfig.
441+
*
440442
* Supports formats:
441443
* - "ssh <host>" or "ssh <user@host>" -> SSH runtime
442444
* - "docker <image>" -> Docker container runtime
@@ -448,57 +450,45 @@ export function parseRuntimeString(
448450
runtime: string | undefined,
449451
_workspaceName: string
450452
): RuntimeConfig | undefined {
451-
if (!runtime) {
452-
return undefined; // Default to worktree (backend decides)
453-
}
454-
455-
const trimmed = runtime.trim();
456-
const lowerTrimmed = trimmed.toLowerCase();
457-
458-
// Worktree runtime (explicit or default)
459-
if (lowerTrimmed === RUNTIME_MODE.WORKTREE) {
460-
return undefined; // Explicit worktree - let backend use default
461-
}
462-
463-
// Local runtime (project-dir, no isolation)
464-
if (lowerTrimmed === RUNTIME_MODE.LOCAL) {
465-
// Return "local" type without srcBaseDir to indicate project-dir runtime
466-
return { type: RUNTIME_MODE.LOCAL };
467-
}
468-
469-
// Parse "ssh <host>" or "ssh <user@host>" format
470-
if (lowerTrimmed === RUNTIME_MODE.SSH || lowerTrimmed.startsWith(SSH_RUNTIME_PREFIX)) {
471-
const hostPart = trimmed.slice(SSH_RUNTIME_PREFIX.length - 1).trim(); // Preserve original case for host
472-
if (!hostPart) {
453+
// Use shared parser from common/types/runtime
454+
const parsed = parseRuntimeModeAndHost(runtime);
455+
456+
// null means invalid input (e.g., "ssh" without host, "docker" without image)
457+
if (parsed === null) {
458+
// Determine which error to throw based on input
459+
const trimmed = runtime?.trim().toLowerCase() ?? "";
460+
if (trimmed === RUNTIME_MODE.SSH || trimmed.startsWith("ssh ")) {
473461
throw new Error("SSH runtime requires host (e.g., 'ssh hostname' or 'ssh user@host')");
474462
}
475-
476-
// Accept both "hostname" and "user@hostname" formats
477-
// SSH will use current user or ~/.ssh/config if user not specified
478-
// Use tilde path - backend will resolve it via runtime.resolvePath()
479-
return {
480-
type: RUNTIME_MODE.SSH,
481-
host: hostPart,
482-
srcBaseDir: "~/mux", // Default remote base directory (tilde will be resolved by backend)
483-
};
484-
}
485-
486-
// Parse "docker <image>" format
487-
if (lowerTrimmed === RUNTIME_MODE.DOCKER || lowerTrimmed.startsWith(DOCKER_RUNTIME_PREFIX)) {
488-
const imagePart = trimmed.slice(DOCKER_RUNTIME_PREFIX.length - 1).trim(); // Preserve original case for image
489-
if (!imagePart) {
463+
if (trimmed === RUNTIME_MODE.DOCKER || trimmed.startsWith("docker ")) {
490464
throw new Error("Docker runtime requires image (e.g., 'docker ubuntu:22.04')");
491465
}
492-
493-
return {
494-
type: RUNTIME_MODE.DOCKER,
495-
image: imagePart,
496-
};
466+
throw new Error(
467+
`Unknown runtime type: '${runtime ?? ""}'. Use 'ssh <host>', 'docker <image>', 'worktree', or 'local'`
468+
);
497469
}
498470

499-
throw new Error(
500-
`Unknown runtime type: '${runtime}'. Use 'ssh <host>', 'docker <image>', 'worktree', or 'local'`
501-
);
471+
// Convert ParsedRuntime to RuntimeConfig
472+
switch (parsed.mode) {
473+
case RUNTIME_MODE.WORKTREE:
474+
return undefined; // Let backend use default worktree config
475+
476+
case RUNTIME_MODE.LOCAL:
477+
return { type: RUNTIME_MODE.LOCAL };
478+
479+
case RUNTIME_MODE.SSH:
480+
return {
481+
type: RUNTIME_MODE.SSH,
482+
host: parsed.host,
483+
srcBaseDir: "~/mux", // Default remote base directory (tilde resolved by backend)
484+
};
485+
486+
case RUNTIME_MODE.DOCKER:
487+
return {
488+
type: RUNTIME_MODE.DOCKER,
489+
image: parsed.image,
490+
};
491+
}
502492
}
503493

504494
export interface CreateWorkspaceOptions {

src/node/runtime/DockerRuntime.ts

Lines changed: 4 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -28,29 +28,19 @@ import type {
2828
import { RuntimeError } from "./Runtime";
2929
import { EXIT_CODE_ABORTED, EXIT_CODE_TIMEOUT } from "@/common/constants/exitCodes";
3030
import { log } from "@/node/services/log";
31-
import { checkInitHookExists, createLineBufferedLoggers, getMuxEnv } from "./initHook";
31+
import { checkInitHookExists, getMuxEnv, runInitHookOnRuntime } from "./initHook";
3232
import { NON_INTERACTIVE_ENV_VARS } from "@/common/constants/env";
3333
import { getProjectName } from "@/node/utils/runtime/helpers";
3434
import { getErrorMessage } from "@/common/utils/errors";
3535
import { DisposableProcess } from "@/node/utils/disposableExec";
36+
import { streamToString, shescape } from "./streamUtils";
3637

3738
/** Hardcoded source directory inside container */
3839
const CONTAINER_SRC_DIR = "/src";
3940

4041
/** Hardcoded background output directory inside container */
4142
const _CONTAINER_BG_OUTPUT_DIR = "/tmp/mux-bashes";
4243

43-
/**
44-
* Shell-escape helper for container bash.
45-
*/
46-
const shescape = {
47-
quote(value: unknown): string {
48-
const s = String(value);
49-
if (s.length === 0) return "''";
50-
return "'" + s.replace(/'/g, "'\"'\"'") + "'";
51-
},
52-
};
53-
5444
/**
5545
* Result of running a Docker command
5646
*/
@@ -568,7 +558,8 @@ export class DockerRuntime implements Runtime {
568558
const hookExists = await checkInitHookExists(projectPath);
569559
if (hookExists) {
570560
const muxEnv = getMuxEnv(projectPath, "docker", branchName);
571-
await this.runInitHook(workspacePath, muxEnv, initLogger, abortSignal);
561+
const hookPath = `${workspacePath}/.mux/init`;
562+
await runInitHookOnRuntime(this, hookPath, workspacePath, muxEnv, initLogger, abortSignal);
572563
} else {
573564
initLogger.logComplete(0);
574565
}
@@ -724,61 +715,6 @@ export class DockerRuntime implements Runtime {
724715
}
725716
}
726717

727-
/**
728-
* Run .mux/init hook inside container
729-
*/
730-
private async runInitHook(
731-
workspacePath: string,
732-
muxEnv: Record<string, string>,
733-
initLogger: InitLogger,
734-
abortSignal?: AbortSignal
735-
): Promise<void> {
736-
const hookPath = `${workspacePath}/.mux/init`;
737-
initLogger.logStep(`Running init hook: ${hookPath}`);
738-
739-
const hookStream = await this.exec(hookPath, {
740-
cwd: workspacePath,
741-
timeout: 3600,
742-
abortSignal,
743-
env: muxEnv,
744-
});
745-
746-
const loggers = createLineBufferedLoggers(initLogger);
747-
const stdoutReader = hookStream.stdout.getReader();
748-
const stderrReader = hookStream.stderr.getReader();
749-
const decoder = new TextDecoder();
750-
751-
const readStdout = async () => {
752-
try {
753-
while (true) {
754-
const { done, value } = await stdoutReader.read();
755-
if (done) break;
756-
loggers.stdout.append(decoder.decode(value, { stream: true }));
757-
}
758-
loggers.stdout.flush();
759-
} finally {
760-
stdoutReader.releaseLock();
761-
}
762-
};
763-
764-
const readStderr = async () => {
765-
try {
766-
while (true) {
767-
const { done, value } = await stderrReader.read();
768-
if (done) break;
769-
loggers.stderr.append(decoder.decode(value, { stream: true }));
770-
}
771-
loggers.stderr.flush();
772-
} finally {
773-
stderrReader.releaseLock();
774-
}
775-
};
776-
777-
const [exitCode] = await Promise.all([hookStream.exitCode, readStdout(), readStderr()]);
778-
779-
initLogger.logComplete(exitCode);
780-
}
781-
782718
// eslint-disable-next-line @typescript-eslint/require-await
783719
async renameWorkspace(
784720
_projectPath: string,
@@ -877,23 +813,3 @@ export class DockerRuntime implements Runtime {
877813
return Promise.resolve("/tmp");
878814
}
879815
}
880-
/**
881-
* Helper to convert a ReadableStream to a string
882-
*/
883-
async function streamToString(stream: ReadableStream<Uint8Array>): Promise<string> {
884-
const reader = stream.getReader();
885-
const decoder = new TextDecoder("utf-8");
886-
let result = "";
887-
888-
try {
889-
while (true) {
890-
const { done, value } = await reader.read();
891-
if (done) break;
892-
result += decoder.decode(value, { stream: true });
893-
}
894-
result += decoder.decode();
895-
return result;
896-
} finally {
897-
reader.releaseLock();
898-
}
899-
}

src/node/runtime/SSHRuntime.ts

Lines changed: 9 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@ import type {
1717
import { RuntimeError as RuntimeErrorClass } from "./Runtime";
1818
import { EXIT_CODE_ABORTED, EXIT_CODE_TIMEOUT } from "@/common/constants/exitCodes";
1919
import { log } from "@/node/services/log";
20-
import { checkInitHookExists, createLineBufferedLoggers, getMuxEnv } from "./initHook";
20+
import { checkInitHookExists, getMuxEnv, runInitHookOnRuntime } from "./initHook";
21+
import { expandTildeForSSH as expandHookPath } from "./tildeExpansion";
2122
import { NON_INTERACTIVE_ENV_VARS } from "@/common/constants/env";
2223
import { streamProcessToLogger } from "./streamProcess";
2324
import { expandTildeForSSH, cdCommandForSSH } from "./tildeExpansion";
@@ -26,20 +27,7 @@ import { getErrorMessage } from "@/common/utils/errors";
2627
import { execAsync, DisposableProcess } from "@/node/utils/disposableExec";
2728
import { getControlPath, sshConnectionPool, type SSHRuntimeConfig } from "./sshConnectionPool";
2829
import { getBashPath } from "@/node/utils/main/bashPath";
29-
30-
/**
31-
* Shell-escape helper for remote bash.
32-
* Reused across all SSH runtime operations for performance.
33-
* Note: For background process commands, use shellQuote from backgroundCommands for parity.
34-
*/
35-
const shescape = {
36-
quote(value: unknown): string {
37-
const s = String(value);
38-
if (s.length === 0) return "''";
39-
// Use POSIX-safe pattern to embed single quotes within single-quoted strings
40-
return "'" + s.replace(/'/g, "'\"'\"'") + "'";
41-
},
42-
};
30+
import { streamToString, shescape } from "./streamUtils";
4331

4432
// Re-export SSHRuntimeConfig from connection pool (defined there to avoid circular deps)
4533
export type { SSHRuntimeConfig } from "./sshConnectionPool";
@@ -766,78 +754,6 @@ export class SSHRuntime implements Runtime {
766754
}
767755
}
768756

769-
/**
770-
* Run .mux/init hook on remote machine if it exists
771-
* @param workspacePath - Path to the workspace directory on remote
772-
* @param muxEnv - MUX_ environment variables (from getMuxEnv)
773-
* @param initLogger - Logger for streaming output
774-
* @param abortSignal - Optional abort signal
775-
*/
776-
private async runInitHook(
777-
workspacePath: string,
778-
muxEnv: Record<string, string>,
779-
initLogger: InitLogger,
780-
abortSignal?: AbortSignal
781-
): Promise<void> {
782-
// Construct hook path - expand tilde if present
783-
const remoteHookPath = `${workspacePath}/.mux/init`;
784-
initLogger.logStep(`Running init hook: ${remoteHookPath}`);
785-
786-
// Expand tilde in hook path for execution
787-
// Tilde won't be expanded when the path is quoted, so we need to expand it ourselves
788-
const hookCommand = expandTildeForSSH(remoteHookPath);
789-
790-
// Run hook remotely and stream output
791-
// No timeout - user init hooks can be arbitrarily long
792-
const hookStream = await this.exec(hookCommand, {
793-
cwd: workspacePath, // Run in the workspace directory
794-
timeout: 3600, // 1 hour - generous timeout for init hooks
795-
abortSignal,
796-
env: muxEnv,
797-
});
798-
799-
// Create line-buffered loggers
800-
const loggers = createLineBufferedLoggers(initLogger);
801-
802-
// Stream stdout/stderr through line-buffered loggers
803-
const stdoutReader = hookStream.stdout.getReader();
804-
const stderrReader = hookStream.stderr.getReader();
805-
const decoder = new TextDecoder();
806-
807-
// Read stdout in parallel
808-
const readStdout = async () => {
809-
try {
810-
while (true) {
811-
const { done, value } = await stdoutReader.read();
812-
if (done) break;
813-
loggers.stdout.append(decoder.decode(value, { stream: true }));
814-
}
815-
loggers.stdout.flush();
816-
} finally {
817-
stdoutReader.releaseLock();
818-
}
819-
};
820-
821-
// Read stderr in parallel
822-
const readStderr = async () => {
823-
try {
824-
while (true) {
825-
const { done, value } = await stderrReader.read();
826-
if (done) break;
827-
loggers.stderr.append(decoder.decode(value, { stream: true }));
828-
}
829-
loggers.stderr.flush();
830-
} finally {
831-
stderrReader.releaseLock();
832-
}
833-
};
834-
835-
// Wait for completion
836-
const [exitCode] = await Promise.all([hookStream.exitCode, readStdout(), readStderr()]);
837-
838-
initLogger.logComplete(exitCode);
839-
}
840-
841757
getWorkspacePath(projectPath: string, workspaceName: string): string {
842758
const projectName = getProjectName(projectPath);
843759
return path.posix.join(this.config.srcBaseDir, projectName, workspaceName);
@@ -952,11 +868,13 @@ export class SSHRuntime implements Runtime {
952868
await this.pullLatestFromOrigin(workspacePath, trunkBranch, initLogger, abortSignal);
953869

954870
// 4. Run .mux/init hook if it exists
955-
// Note: runInitHook calls logComplete() internally if hook exists
871+
// Note: runInitHookOnRuntime calls logComplete() internally
956872
const hookExists = await checkInitHookExists(projectPath);
957873
if (hookExists) {
958874
const muxEnv = getMuxEnv(projectPath, "ssh", branchName);
959-
await this.runInitHook(workspacePath, muxEnv, initLogger, abortSignal);
875+
// Expand tilde in hook path (quoted paths don't auto-expand on remote)
876+
const hookPath = expandHookPath(`${workspacePath}/.mux/init`);
877+
await runInitHookOnRuntime(this, hookPath, workspacePath, muxEnv, initLogger, abortSignal);
960878
} else {
961879
// No hook - signal completion immediately
962880
initLogger.logComplete(0);
@@ -1325,23 +1243,5 @@ export class SSHRuntime implements Runtime {
13251243
}
13261244
}
13271245

1328-
/**
1329-
* Helper to convert a ReadableStream to a string
1330-
*/
1331-
export async function streamToString(stream: ReadableStream<Uint8Array>): Promise<string> {
1332-
const reader = stream.getReader();
1333-
const decoder = new TextDecoder("utf-8");
1334-
let result = "";
1335-
1336-
try {
1337-
while (true) {
1338-
const { done, value } = await reader.read();
1339-
if (done) break;
1340-
result += decoder.decode(value, { stream: true });
1341-
}
1342-
result += decoder.decode();
1343-
return result;
1344-
} finally {
1345-
reader.releaseLock();
1346-
}
1347-
}
1246+
// Re-export for backward compatibility with existing imports
1247+
export { streamToString } from "./streamUtils";

0 commit comments

Comments
 (0)