Fix enclave smoke runtime and rootless cleanup failures

src/artifact-permissions.test.ts

@@ -156,4 +156,43 @@ describe('artifact-permissions', () => {
       fs.rmSync(auditDir, { recursive: true, force: true });
     }
   });
+
+  it('uses an explicitly available local image for rootless permission repair', () => {
+    const auditDir = makeTempDir();
+    try {
+      getuidSpy = jest.spyOn(process, 'getuid').mockReturnValue(1001);
+      const repaired = fixArtifactPermissionsForRootless(
+        [auditDir],
+        undefined,
+        'ghcr.io/github/gh-aw-firewall',
+        'v1@sha256:deadbeef',
+        undefined,
+        'awf-enclave-mcp-server:local',
+      );
+      expect(repaired).toBe(true);
+      expect(mockExecaSync).toHaveBeenCalledWith(
+        'docker',
+        expect.arrayContaining(['awf-enclave-mcp-server:local']),
+        expect.objectContaining({ reject: false }),
+      );
+    } finally {
+      fs.rmSync(auditDir, { recursive: true, force: true });
+    }
+  });
+
+  it('reports failed rootless permission repair', () => {
+    const auditDir = makeTempDir();
+    let warnSpy: jest.SpyInstance | undefined;
+    try {
+      getuidSpy = jest.spyOn(process, 'getuid').mockReturnValue(1001);
+      warnSpy = jest.spyOn(console, 'error').mockImplementation(() => {});
+      mockExecaSync.mockReturnValue({ stdout: '', stderr: 'image unavailable', exitCode: 1 });
+      expect(
+        fixArtifactPermissionsForRootless([auditDir], undefined, undefined, undefined, undefined),
+      ).toBe(false);
+    } finally {
+      warnSpy?.mockRestore();
+      fs.rmSync(auditDir, { recursive: true, force: true });
+    }
+  });
 });

src/artifact-permissions.ts

@@ -60,22 +60,24 @@
   imageRegistry: string | undefined,
   imageTag: string | undefined,
   agentImage: string | undefined,
-): void {
+  imageRefOverride?: string,
+): boolean {
   const currentUid = process.getuid?.();
   if (currentUid === undefined || currentUid === 0) {
-    return;
+    return true;
   }
 
   const existingDirs = dirs.filter(
     (dir): dir is string => typeof dir === 'string' && dir.length > 0 && fs.existsSync(dir),
   );
   if (existingDirs.length === 0) {
-    return;
+    return true;
   }
 
   const uid = getSafeHostUid();
   const gid = getSafeHostGid();
-  const imageRef = resolvePermFixerImageRef(imageRegistry, imageTag, agentImage);
+  const imageRef = imageRefOverride || resolvePermFixerImageRef(imageRegistry, imageTag, agentImage);
+  let repairedAll = true;
 
   for (const dir of existingDirs) {
     const mount = applyHostPathPrefixToVolumes([`${path.resolve(dir)}:/fix:rw`], dockerHostPathPrefix)[0];
@@ -100,9 +102,9 @@
           '--entrypoint',
           'sh',
           '-e',
           `TUID=${uid}`,
           '-e',
           `TGID=${gid}`,
           '-v',
           mount,
           imageRef,
@@ -132,9 +134,12 @@
         } else {
           logger.warn(`Rootless artifact permission repair failed ${detail}`);
         }
+        repairedAll = false;
       }
     } catch (error) {
       logger.warn(`Rootless artifact permission repair failed for ${dir}:`, error);
+      repairedAll = false;
     }
   }
+  return repairedAll;
 }

src/constants.ts

@@ -11,6 +11,7 @@ export const DOH_PROXY_CONTAINER_NAME = 'awf-doh-proxy';
 export const CLI_PROXY_CONTAINER_NAME = 'awf-cli-proxy';
 export const ENCLAVE_MCP_SERVER_CONTAINER_NAME = 'awf-enclave-mcp-server';
 export const ENCLAVE_AGENT_API_PROXY_CONTAINER_NAME = 'awf-enclave-agent-api-proxy';
+export const LOCAL_ENCLAVE_MCP_SERVER_IMAGE = 'awf-enclave-mcp-server:local';
 
 // SQUID_PORT is centralized in src/config/sandbox-network-policy.json and
 // re-exported here so existing import sites keep working unchanged.

src/enclave/gateway.test.ts

@@ -537,7 +537,7 @@ describe('enclave mcpg handoff', () => {
   it('drains the AWF server and disconnects mcpg without stopping the external container', async () => {
     mockExeca
       .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' })
-      .mockResolvedValueOnce({ exitCode: 0, stdout: '0\n', stderr: '' })
+      .mockResolvedValueOnce({ exitCode: 0, stdout: '0|false|\n', stderr: '' })
       .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' });
     await shutdownEnclaveGateway(config(), env());
     expect(mockExeca).toHaveBeenNthCalledWith(
@@ -549,7 +549,11 @@ describe('enclave mcpg handoff', () => {
     expect(mockExeca).toHaveBeenNthCalledWith(
       2,
       'docker',
-      ['inspect', '--format={{.State.ExitCode}}', 'awf-enclave-mcp-server'],
+      [
+        'inspect',
+        '--format={{.State.ExitCode}}|{{.State.OOMKilled}}|{{.State.Error}}',
+        'awf-enclave-mcp-server',
+      ],
       expect.anything(),
     );
     expect(mockExeca).toHaveBeenNthCalledWith(
@@ -566,4 +570,13 @@ describe('enclave mcpg handoff', () => {
     await expect(shutdownEnclaveGateway(config(), env())).rejects.toThrow(/Failed to drain/);
     expect(mockExeca).toHaveBeenCalledTimes(1);
   });
+
+  it('reports enclave server OOM state after an abnormal shutdown', async () => {
+    mockExeca
+      .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' })
+      .mockResolvedValueOnce({ exitCode: 0, stdout: '137|true|\n', stderr: '' });
+    await expect(shutdownEnclaveGateway(config(), env()))
+      .rejects.toThrow(/137\|true\|/);
+    expect(mockExeca).toHaveBeenCalledTimes(2);
+  });
 });

src/enclave/gateway.ts

@@ -523,11 +523,18 @@ export async function shutdownEnclaveGateway(
   }
   const inspectResult = await execa(
     'docker',
-    ['inspect', '--format={{.State.ExitCode}}', ENCLAVE_MCP_SERVER_CONTAINER_NAME],
+    [
+      'inspect',
+      '--format={{.State.ExitCode}}|{{.State.OOMKilled}}|{{.State.Error}}',
+      ENCLAVE_MCP_SERVER_CONTAINER_NAME,
+    ],
     { env: getLocalDockerEnv(), reject: false, timeout: 10_000 },
   );
-  if (inspectResult.exitCode !== 0 || inspectResult.stdout.trim() !== '0') {
-    throw new Error('Enclave MCP server did not complete graceful cleanup');
+  const state = inspectResult.stdout.trim();
+  if (inspectResult.exitCode !== 0 || !state.startsWith('0|false|')) {
+    throw new Error(
+      `Enclave MCP server did not complete graceful cleanup (state: ${state || 'unavailable'})`,
+    );
   }
   await execa(
     'docker',

src/enclave/manager.test.ts

@@ -5,6 +5,7 @@ import execa from 'execa';
 import { normalizeEnclavesConfig } from '../parsers/enclave-parser';
 import type { WrapperConfig } from '../types';
 import {
+  enclaveManagerTestHelpers,
   isEnclaveAgentEnabled,
   isEnclaveScriptEnabled,
   prepareEnclaves,
@@ -35,6 +36,7 @@ function enclaveEnv(overrides: NodeJS.ProcessEnv = {}): NodeJS.ProcessEnv {
     AWF_ENCLAVE_MCP_CAPABILITY: 'a'.repeat(64),
     AWF_ENCLAVE_MCP_GATEWAY_IDENTITY: 'test-run-identity',
     AWF_ENCLAVE_MCP_GATEWAY_ENDPOINT: 'http://127.0.0.1:8080/mcp/awf-enclave',
+    MCP_GATEWAY_API_KEY: 'g'.repeat(48),
     ...overrides,
   };
 }
@@ -305,4 +307,33 @@ describe('prepareEnclaves fail-closed preflight', () => {
     expect(fs.existsSync(paths.root)).toBe(true);
     expect(fs.existsSync(paths.ingressRoot)).toBe(true);
   });
+
+  it('uses the local enclave image and defers cleanup when permission repair fails', async () => {
+    const wrapperConfig = config(workDir);
+    wrapperConfig.buildLocal = true;
+    const paths = resolveEnclavePaths(workDir);
+    fs.mkdirSync(paths.root, { recursive: true });
+    fs.mkdirSync(paths.ingressRoot, { recursive: true });
+    const remove = jest.fn((target: string) => {
+      if (target === paths.root) {
+        throw Object.assign(new Error('permission denied'), { code: 'EACCES' });
+      }
+    });
+    const repair = jest.fn(() => false);
+    const warnSpy = jest.spyOn(console, 'error').mockImplementation(() => {});
+    try {
+      enclaveManagerTestHelpers.removePrivateState(wrapperConfig, paths, { remove, repair });
+      expect(repair).toHaveBeenCalledWith(
+        [paths.root, paths.ingressRoot],
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+        'awf-enclave-mcp-server:local',
+      );
+      expect(remove.mock.calls.filter(([target]) => target === paths.root)).toHaveLength(1);
+    } finally {
+      warnSpy.mockRestore();
+    }
+  });
 });

src/enclave/manager.ts

@@ -10,6 +10,7 @@ import { releaseSeedPermissions, resolveStagingToken, stageEnclaveSeeds, type Gi
 import { getLocalDockerEnv } from '../host-env';
 import { getSafeHostGid, getSafeHostUid } from '../host-identity';
 import { logger } from '../logger';
+import { LOCAL_ENCLAVE_MCP_SERVER_IMAGE } from '../constants';
 import type { WrapperConfig } from '../types';
 import type {
   EnclaveAgentExecutorConfig,
@@ -220,21 +221,38 @@ async function removeOrphanEnclaveContainers(runId: string): Promise<void> {
   }
 }
 
-function removePrivateState(config: WrapperConfig, paths: EnclavePaths): void {
+interface RemovePrivateStateDependencies {
+  remove?: (target: string) => void;
+  repair?: typeof fixArtifactPermissionsForRootless;
+}
+
+function removePrivateState(
+  config: WrapperConfig,
+  paths: EnclavePaths,
+  dependencies: RemovePrivateStateDependencies = {},
+): void {
+  const remove = dependencies.remove
+    ?? ((target: string) => fs.rmSync(target, { recursive: true, force: true }));
+  const repair = dependencies.repair ?? fixArtifactPermissionsForRootless;
   try {
-    fs.rmSync(paths.root, { recursive: true, force: true });
-    fs.rmSync(paths.ingressRoot, { recursive: true, force: true });
+    remove(paths.root);
+    remove(paths.ingressRoot);
   } catch (error: unknown) {
     if (error && typeof error === 'object' && 'code' in error && error.code === 'EACCES') {
-      fixArtifactPermissionsForRootless(
+      const repaired = repair(
         [paths.root, paths.ingressRoot],
         config.dockerHostPathPrefix,
         config.imageRegistry,
         config.imageTag,
         config.agentImage,
+        config.buildLocal ? LOCAL_ENCLAVE_MCP_SERVER_IMAGE : undefined,
       );
-      fs.rmSync(paths.root, { recursive: true, force: true });
-      fs.rmSync(paths.ingressRoot, { recursive: true, force: true });
+      if (!repaired) {
+        logger.warn('Enclaves: private state permission repair failed; cleanup will be deferred');
+        return;
+      }
+      remove(paths.root);
+      remove(paths.ingressRoot);
       return;
     }
     throw error;
@@ -264,5 +282,6 @@ export async function teardownEnclaves(config: WrapperConfig): Promise<void> {
 export const enclaveManagerTestHelpers = {
   prepareDirectories,
   readRunId,
+  removePrivateState,
   removeOrphanEnclaveContainers,
 };

src/services/enclave-mcp-service.test.ts

@@ -42,6 +42,8 @@ describe('buildEnclaveMcpService', () => {
     expect(result.service).toMatchObject({
       container_name: 'awf-enclave-mcp-server',
       image: 'ghcr.io/github/gh-aw-firewall/enclave-mcp-server:v1',
+      mem_limit: '1g',
+      memswap_limit: '1g',
       depends_on: {
         'enclave-script-image': { condition: 'service_completed_successfully' },
       },

src/services/enclave-mcp-service.ts

@@ -2,6 +2,7 @@ import { buildRuntimeImageRef } from '../image-tag';
 import {
   ENCLAVE_AGENT_API_PROXY_CONTAINER_NAME,
   ENCLAVE_MCP_SERVER_CONTAINER_NAME,
+  LOCAL_ENCLAVE_MCP_SERVER_IMAGE,
 } from '../constants';
 import type { WrapperConfig } from '../types';
 import { API_PROXY_PORTS } from '../types/ports';
@@ -64,7 +65,6 @@ import {
 
 const LOCAL_ENCLAVE_SCRIPT_IMAGE = 'awf-enclave-script:local';
 const LOCAL_ENCLAVE_AGENT_IMAGE = 'awf-enclave-agent:local';
-const LOCAL_ENCLAVE_MCP_SERVER_IMAGE = 'awf-enclave-mcp-server:local';
 const ENCLAVE_SCRIPT_IMAGE_NAME = 'enclave-script';
 const ENCLAVE_AGENT_IMAGE_NAME = 'enclave-agent';
 const ENCLAVE_MCP_SERVER_IMAGE_NAME = 'enclave-mcp-server';
@@ -377,7 +377,7 @@ export function buildEnclaveMcpService(params: EnclaveMcpServiceParams): Enclave
       retries: 10,
       start_period: '20s',
     },
-    ...buildContainerSecurityHardening({ memLimit: '256m', pidsLimit: 100, cpuShares: 256 }),
+    ...buildContainerSecurityHardening({ memLimit: '1g', pidsLimit: 100, cpuShares: 256 }),
     cap_add: ['CHOWN', 'DAC_OVERRIDE', 'FOWNER'],
     restart: 'no',
     stop_grace_period: '5s',