🤖 fix: reuse cached ssh layouts across runtime instances

Author: ammar-agent

src/node/runtime/SSHRuntime.test.ts

@@ -4,7 +4,7 @@ import { describe, expect, it, beforeEach, afterEach, spyOn } from "bun:test";
 import * as runtimeHelpers from "@/node/utils/runtime/helpers";
 import * as disposableExec from "@/node/utils/disposableExec";
 import * as submoduleSync from "./submoduleSync";
-import { SSHRuntime, computeBaseRepoPath } from "./SSHRuntime";
+import { SSHRuntime, clearSharedProjectLayoutCache, computeBaseRepoPath } from "./SSHRuntime";
 import {
   buildLegacyRemoteProjectLayout,
   buildRemoteProjectLayout,
@@ -35,6 +35,10 @@ function createMockExecResult(
   } as unknown as ReturnType<typeof disposableExec.execFileAsync>;
 }
 
+afterEach(() => {
+  clearSharedProjectLayoutCache();
+});
+
 describe("SSHRuntime constructor", () => {
   it("should accept tilde in srcBaseDir", () => {
     // Tildes are now allowed - they will be resolved via resolvePath()
@@ -1205,6 +1209,23 @@ describe("SSHRuntime layout detection", () => {
     expect(detectionCommand).toContain(`test -e "${legacyLayout.projectRoot}/${workspaceName}"`);
     expect(detectionCommand).not.toContain(`test -d "${legacyLayout.projectRoot}"`);
   });
+  it("reuses a cached legacy layout for fresh runtimes without workspacePath hints", () => {
+    const config = { host: "example.com", srcBaseDir: "/home/user/src" };
+    const projectPath = "/projects/cached-legacy-demo";
+    const workspaceName = "legacy-slot";
+    const legacyLayout = buildLegacyRemoteProjectLayout(config.srcBaseDir, projectPath);
+    const legacyWorkspacePath = getRemoteWorkspacePath(legacyLayout, workspaceName);
+
+    new SSHRuntime(config, createSSHTransport(config, false), {
+      projectPath,
+      workspaceName,
+      workspacePath: legacyWorkspacePath,
+    });
+
+    const freshRuntime = new SSHRuntime(config, createSSHTransport(config, false));
+
+    expect(freshRuntime.getWorkspacePath(projectPath, workspaceName)).toBe(legacyWorkspacePath);
+  });
 });
 
 describe("computeBaseRepoPath", () => {

src/node/runtime/SSHRuntime.ts

@@ -62,6 +62,18 @@ const BUNDLE_REF_PREFIX = "refs/mux-bundle/";
 /** Small backoff for concurrent writers healing the same shared base repo config. */
 const BASE_REPO_CONFIG_LOCK_RETRY_DELAYS_MS = [50, 100, 200];
 
+const sharedProjectLayouts = new Map<string, RemoteProjectLayout>();
+
+function getProjectLayoutCacheKey(config: SSHRuntimeConfig, projectPath: string): string {
+  return [
+    config.host,
+    config.port?.toString() ?? "22",
+    config.identityFile ?? "default",
+    config.srcBaseDir,
+    projectPath,
+  ].join(":");
+}
+
 function isGitConfigLockConflict(message: string): boolean {
   return /could not lock config file/i.test(message);
 }
@@ -166,6 +178,10 @@ export function computeBaseRepoPath(srcBaseDir: string, projectPath: string): st
  *
  * Extends RemoteRuntime for shared exec/file operations.
  */
+export function clearSharedProjectLayoutCache(): void {
+  sharedProjectLayouts.clear();
+}
+
 export class SSHRuntime extends RemoteRuntime {
   private readonly config: SSHRuntimeConfig;
   private readonly transport: SSHTransport;
@@ -195,7 +211,7 @@ export class SSHRuntime extends RemoteRuntime {
     this.currentWorkspacePath = options?.workspacePath;
 
     if (options?.projectPath && options.workspacePath) {
-      this.projectLayouts.set(
+      this.cacheProjectLayout(
         options.projectPath,
         buildRemoteProjectLayout(
           this.config.srcBaseDir,
@@ -250,15 +266,31 @@ export class SSHRuntime extends RemoteRuntime {
     return buildRemoteProjectLayout(this.config.srcBaseDir, projectPath);
   }
 
+  private getCachedProjectLayout(projectPath: string): RemoteProjectLayout | undefined {
+    return (
+      this.projectLayouts.get(projectPath) ??
+      sharedProjectLayouts.get(getProjectLayoutCacheKey(this.config, projectPath))
+    );
+  }
+
+  private cacheProjectLayout(
+    projectPath: string,
+    layout: RemoteProjectLayout
+  ): RemoteProjectLayout {
+    this.projectLayouts.set(projectPath, layout);
+    sharedProjectLayouts.set(getProjectLayoutCacheKey(this.config, projectPath), layout);
+    return layout;
+  }
+
   private getPreferredProjectLayout(projectPath: string): RemoteProjectLayout {
-    return this.projectLayouts.get(projectPath) ?? this.getDefaultProjectLayout(projectPath);
+    return this.getCachedProjectLayout(projectPath) ?? this.getDefaultProjectLayout(projectPath);
   }
 
   private async resolveProjectLayout(
     projectPath: string,
     workspaceName?: string
   ): Promise<RemoteProjectLayout> {
-    const cached = this.projectLayouts.get(projectPath);
+    const cached = this.getCachedProjectLayout(projectPath);
     if (cached) {
       return cached;
     }
@@ -286,11 +318,9 @@ export class SSHRuntime extends RemoteRuntime {
         timeout: 10,
       });
       const layout = detection.stdout.trim() === "legacy" ? legacyLayout : preferredLayout;
-      this.projectLayouts.set(projectPath, layout);
-      return layout;
+      return this.cacheProjectLayout(projectPath, layout);
     } catch {
-      this.projectLayouts.set(projectPath, preferredLayout);
-      return preferredLayout;
+      return this.cacheProjectLayout(projectPath, preferredLayout);
     }
   }
 
@@ -435,12 +465,7 @@ export class SSHRuntime extends RemoteRuntime {
       return this.currentWorkspacePath;
     }
 
-    const cachedLayout = this.projectLayouts.get(projectPath);
-    if (cachedLayout) {
-      return getRemoteWorkspacePath(cachedLayout, workspaceName);
-    }
-
-    return getRemoteWorkspacePath(this.getDefaultProjectLayout(projectPath), workspaceName);
+    return getRemoteWorkspacePath(this.getPreferredProjectLayout(projectPath), workspaceName);
   }
 
   /**

src/node/services/utils/forkOrchestrator.test.ts

@@ -153,6 +153,7 @@ describe("orchestrateFork", () => {
     expect(createRuntimeMock).toHaveBeenCalledWith(DEFAULT_FORKED_RUNTIME_CONFIG, {
       projectPath: PROJECT_PATH,
       workspaceName: NEW_WORKSPACE_NAME,
+      workspacePath: "/workspaces/forked",
     });
   });
 
@@ -201,6 +202,7 @@ describe("orchestrateFork", () => {
     expect(createRuntimeMock).toHaveBeenCalledWith(DEFAULT_FORKED_RUNTIME_CONFIG, {
       projectPath: PROJECT_PATH,
       workspaceName: NEW_WORKSPACE_NAME,
+      workspacePath: "/workspaces/created",
     });
   });
 
@@ -409,6 +411,7 @@ describe("orchestrateFork", () => {
     expect(createRuntimeMock).toHaveBeenCalledWith(customForkedRuntimeConfig, {
       projectPath: PROJECT_PATH,
       workspaceName: NEW_WORKSPACE_NAME,
+      workspacePath: "/workspaces/created-with-custom-runtime",
     });
   });
 
@@ -451,10 +454,14 @@ describe("orchestrateFork", () => {
       expect.objectContaining({ containerName: "mux-demo-source-aaaaaa" })
     );
 
-    // createRuntime should also receive the normalized config
+    // createRuntime should also receive the normalized config and the created workspace path.
     expect(createRuntimeMock).toHaveBeenCalledWith(
       expect.objectContaining({ containerName: expectedContainerName }),
-      { projectPath: PROJECT_PATH, workspaceName: NEW_WORKSPACE_NAME }
+      {
+        projectPath: PROJECT_PATH,
+        workspaceName: NEW_WORKSPACE_NAME,
+        workspacePath: "/workspaces/new",
+      }
     );
   });
   it("returns Err when create fallback also fails", async () => {

src/node/services/utils/forkOrchestrator.ts

@@ -577,6 +577,7 @@ export async function orchestrateFork(
   const targetRuntime = createRuntime(normalizedForkedRuntimeConfig, {
     projectPath,
     workspaceName: newWorkspaceName,
+    workspacePath,
   });
 
   return Ok({

src/node/services/workspaceService.ts

@@ -5168,9 +5168,11 @@ export class WorkspaceService extends EventEmitter {
       // Copy plan file using explicit source/target runtimes for cross-runtime safety.
       // Create a fresh source runtime handle because DockerRuntime.forkWorkspace() can
       // mutate the original runtime's container identity to target the new workspace.
+      const sourceWorkspace = this.config.findWorkspace(sourceWorkspaceId);
       const freshSourceRuntime = createRuntime(sourceRuntimeConfig, {
         projectPath: foundProjectPath,
         workspaceName: sourceMetadata.name,
+        workspacePath: sourceWorkspace?.workspacePath,
       });
       await copyPlanFileAcrossRuntimes(
         freshSourceRuntime,