fix(cua): isolate caller image rebuilds

Signed-off-by: Julie Yaunches <jyaunches@nvidia.com>

Author: jyaunches

src/lib/actions/sandbox/rebuild-flow-helpers.test.ts

@@ -1,7 +1,12 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+import fs from "node:fs";
+import os from "node:os";
+import path from "node:path";
+
 import { afterEach, beforeEach, describe, expect, it, type MockInstance, vi } from "vitest";
+
 import * as dockerImage from "../../adapters/docker/image";
 import * as agentDefs from "../../agent/defs";
 import * as agentOnboard from "../../agent/onboard";
@@ -189,6 +194,91 @@ describe("rebuild agent base image preflight", () => {
     });
   });
 
+  it("hands a pinned NemoCUA image alias to the inner sandbox create (#9649)", () => {
+    const cuaOverrideEnvVar = "NEMOCLAW_CUA_SANDBOX_IMAGE_REF";
+    const mutableRef = "nemocua-scenario:mutable";
+    const pinnedRef = `nemoclaw-nemocua-sandbox-base-local:rebuild-1-${"a".repeat(16)}-image-${"b".repeat(64)}`;
+    const agent = agentDefs.loadAgent("nemocua", { NEMOCLAW_CUA_ENABLED: "1" });
+    const root = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-cua-rebuild-test-"));
+    let buildContext = root;
+    vi.stubEnv("NEMOCLAW_CUA_ENABLED", "1");
+    vi.stubEnv(cuaOverrideEnvVar, mutableRef);
+    vi.spyOn(agentDefs, "loadAgent").mockReturnValue(agent);
+    vi.spyOn(agentOnboard, "ensureAgentBaseImage").mockReturnValue({
+      imageTag: mutableRef,
+      built: false,
+    });
+    const pinImage = vi
+      .spyOn(agentOnboard, "pinAgentSandboxBaseImageRef")
+      .mockReturnValue(pinnedRef);
+    const dockerRmi = vi.spyOn(dockerImage, "dockerRmi").mockReturnValue({ status: 0 } as never);
+
+    try {
+      const preflight = ensureRebuildAgentBaseImage("nemocua", makeBail());
+      expect(preflight).toMatchObject({
+        ok: true,
+        imageRef: pinnedRef,
+        overrideEnvVar: cuaOverrideEnvVar,
+      });
+      expect(pinImage).toHaveBeenCalledWith("nemocua", mutableRef, {
+        forceLocal: true,
+        temporary: true,
+      });
+
+      const restore = pinRebuildAgentBaseImageForRecreate(preflight);
+      try {
+        const inner = agentOnboard.createAgentSandbox(agent, { rootDir: root });
+        buildContext = inner.buildCtx;
+        const dockerfile = fs.readFileSync(inner.stagedDockerfile, "utf8");
+        expect(dockerfile).toContain(`ARG BASE_IMAGE=${pinnedRef}`);
+        expect(dockerfile).not.toContain(mutableRef);
+      } finally {
+        restore();
+      }
+
+      expect(process.env[cuaOverrideEnvVar]).toBe(mutableRef);
+      expect(disposeRebuildAgentBaseImagePreflight(preflight)).toBe(true);
+      expect(dockerRmi).toHaveBeenCalledWith(pinnedRef, {
+        ignoreError: true,
+        suppressOutput: true,
+      });
+    } finally {
+      fs.rmSync(buildContext, { recursive: true, force: true });
+      fs.rmSync(root, { recursive: true, force: true });
+      vi.unstubAllEnvs();
+    }
+  });
+
+  it("passes an immutable NemoCUA digest through without a local alias (#9649)", () => {
+    const cuaOverrideEnvVar = "NEMOCLAW_CUA_SANDBOX_IMAGE_REF";
+    const digestRef = `registry.example/nemocua@sha256:${"a".repeat(64)}`;
+    const agent = agentDefs.loadAgent("nemocua", { NEMOCLAW_CUA_ENABLED: "1" });
+    vi.stubEnv("NEMOCLAW_CUA_ENABLED", "1");
+    vi.stubEnv(cuaOverrideEnvVar, digestRef);
+    vi.spyOn(agentDefs, "loadAgent").mockReturnValue(agent);
+    vi.spyOn(agentOnboard, "ensureAgentBaseImage").mockReturnValue({
+      imageTag: digestRef,
+      built: false,
+    });
+    const pinImage = vi.spyOn(agentOnboard, "pinAgentSandboxBaseImageRef");
+    const dockerRmi = vi.spyOn(dockerImage, "dockerRmi");
+
+    try {
+      const preflight = ensureRebuildAgentBaseImage("nemocua", makeBail());
+
+      expect(preflight).toEqual({
+        ok: true,
+        imageRef: digestRef,
+        overrideEnvVar: cuaOverrideEnvVar,
+      });
+      expect(pinImage).not.toHaveBeenCalled();
+      expect(disposeRebuildAgentBaseImagePreflight(preflight)).toBe(true);
+      expect(dockerRmi).not.toHaveBeenCalled();
+    } finally {
+      vi.unstubAllEnvs();
+    }
+  });
+
   it("fails closed when an explicit local result lacks validated outer metadata", () => {
     process.env[overrideEnvVar] = "nemoclaw-hermes-sandbox-base-local:caller";
     const mutableRef = "nemoclaw-hermes-sandbox-base-local:resolved";

src/lib/actions/sandbox/rebuild-flow-helpers.ts

@@ -315,6 +315,22 @@ export function ensureRebuildAgentBaseImage(
     } finally {
       restoreExplicitOverrideTrust();
     }
+    if (agentDef.name === "nemocua") {
+      if (!result.imageTag) throw new Error("NemoCUA caller image resolution returned no image");
+      if (isImmutableRemoteBaseImageRef(result.imageTag)) {
+        return { ok: true, imageRef: result.imageTag, overrideEnvVar };
+      }
+      const imageRef = pinAgentSandboxBaseImageRef(agentDef.name, result.imageTag, {
+        forceLocal: true,
+        temporary: true,
+      });
+      return {
+        ok: true,
+        imageRef,
+        overrideEnvVar,
+        disposeImageRef: createTemporaryBaseImageHandoffDisposer(imageRef),
+      };
+    }
     const reusedLocalResolution =
       result.resolutionMetadata?.source === "local" &&
       result.reusedResolutionHint === result.resolutionMetadata;

src/lib/agent/base-image.test.ts

@@ -108,25 +108,20 @@ describe("agent base image provisioning", () => {
     });
   });
 
-  it("stages NemoCUA through the ordinary agent Dockerfile path (#9649)", () => {
+  it("stages only the NemoCUA Dockerfile in the caller-image build context (#9649)", () => {
     vi.stubEnv("NEMOCLAW_CUA_ENABLED", "1");
     vi.stubEnv("NEMOCLAW_CUA_SANDBOX_IMAGE_REF", "nemocua-scenario:staged");
     const agent = loadAgent("nemocua");
     const root = tmpDir();
-    const stagedAgentDir = path.join(root, "agents", "nemocua");
-    fs.mkdirSync(path.dirname(stagedAgentDir), { recursive: true });
-    fs.cpSync(agent.agentDir, stagedAgentDir, { recursive: true });
+    fs.writeFileSync(path.join(root, "unrelated-sentinel.txt"), "must not enter build context");
     let buildContext = root;
 
     try {
       withMockedDocker(({ createAgentSandbox }) => {
         const result = createAgentSandbox(agent, { rootDir: root });
         buildContext = result.buildCtx;
-        expect(fs.readdirSync(path.join(result.buildCtx, "agents", "nemocua")).sort()).toEqual([
-          "Dockerfile",
-          "manifest.yaml",
-          "policy-additions.yaml",
-        ]);
+        expect(fs.readdirSync(result.buildCtx)).toEqual(["Dockerfile"]);
+        expect(fs.existsSync(path.join(result.buildCtx, "unrelated-sentinel.txt"))).toBe(false);
         expect(fs.readFileSync(result.stagedDockerfile, "utf8")).toContain(
           "ARG BASE_IMAGE=nemocua-scenario:staged",
         );

src/lib/agent/base-image.ts

@@ -14,7 +14,7 @@ import {
   dockerRmi,
   dockerTag,
 } from "../adapters/docker";
-import { requireCuaSandboxImageRef } from "../cua/feature";
+import { CUA_SANDBOX_IMAGE_ENV, requireCuaSandboxImageRef } from "../cua/feature";
 import { encodeCorporateCaArg, resolveCorporateCa } from "../onboard/corporate-ca";
 import { createCustomBuildContextFilter } from "../onboard/custom-build-context";
 import { ROOT } from "../runner";
@@ -172,6 +172,7 @@ function reuseTrustedAgentRemoteBaseImageOverride(
 }
 
 export function getAgentSandboxBaseImageEnvVar(agentName: string): string {
+  if (agentName === "nemocua") return CUA_SANDBOX_IMAGE_ENV;
   return `NEMOCLAW_${agentName.toUpperCase().replace(/[^A-Z0-9]+/g, "_")}_SANDBOX_BASE_IMAGE_REF`;
 }
 
@@ -724,11 +725,13 @@ export function createAgentSandbox(
   const buildCtx = fs.mkdtempSync(path.join(os.tmpdir(), SANDBOX_BUILD_CONTEXT_PREFIX));
   const stagedDockerfile = path.join(buildCtx, "Dockerfile");
   try {
-    const shouldIncludeBuildContextPath = createCustomBuildContextFilter(rootDir);
-    fs.cpSync(rootDir, buildCtx, {
-      recursive: true,
-      filter: (src) => path.basename(src) !== ".claude" && shouldIncludeBuildContextPath(src),
-    });
+    if (agent.name !== "nemocua") {
+      const shouldIncludeBuildContextPath = createCustomBuildContextFilter(rootDir);
+      fs.cpSync(rootDir, buildCtx, {
+        recursive: true,
+        filter: (src) => path.basename(src) !== ".claude" && shouldIncludeBuildContextPath(src),
+      });
+    }
     fs.copyFileSync(agentDockerfile, stagedDockerfile);
     if (baseImageRef) {
       const dockerfile = fs.readFileSync(stagedDockerfile, "utf8");