refactor(config): simplify env snapshot write context

src/config/env-preserve-io.test.ts

@@ -22,6 +22,35 @@ async function withTempConfig(
   }
 }
 
+async function withEnvOverrides(
+  updates: Record<string, string | undefined>,
+  run: () => Promise<void>,
+): Promise<void> {
+  const previous = new Map<string, string | undefined>();
+  for (const key of Object.keys(updates)) {
+    previous.set(key, process.env[key]);
+  }
+
+  try {
+    for (const [key, value] of Object.entries(updates)) {
+      if (value === undefined) {
+        delete process.env[key];
+      } else {
+        process.env[key] = value;
+      }
+    }
+    await run();
+  } finally {
+    for (const [key, value] of previous.entries()) {
+      if (value === undefined) {
+        delete process.env[key];
+      } else {
+        process.env[key] = value;
+      }
+    }
+  }
+}
+
 describe("env snapshot TOCTOU via createConfigIO", () => {
   it("restores env refs using read-time env even after env mutation", async () => {
     const env: Record<string, string> = {
@@ -33,20 +62,17 @@ describe("env snapshot TOCTOU via createConfigIO", () => {
     await withTempConfig(configJson, async (configPath) => {
       // Instance A: read config (captures env snapshot)
       const ioA = createConfigIO({ configPath, env: env as unknown as NodeJS.ProcessEnv });
-      const snapshot = await ioA.readConfigFileSnapshot();
+      const firstRead = await ioA.readConfigFileSnapshotForWrite();
-      expect(snapshot.config.gateway?.remote?.token).toBe("original-key-123");
+      expect(firstRead.snapshot.config.gateway?.remote?.token).toBe("original-key-123");
 
       // Mutate env between read and write
       env.MY_API_KEY = "mutated-key-456";
 
-      // Instance B: write config, but inject snapshot from A
+      // Instance B: write config using explicit read context from A
       const ioB = createConfigIO({ configPath, env: env as unknown as NodeJS.ProcessEnv });
-      const envSnapshot = ioA.getEnvSnapshot();
-      expect(envSnapshot).not.toBeNull();
-      ioB.setEnvSnapshot(envSnapshot!);
 
       // Write the resolved config back — should restore ${MY_API_KEY}
-      await ioB.writeConfigFile(snapshot.config);
+      await ioB.writeConfigFile(firstRead.snapshot.config, firstRead.writeOptions);
 
       // Verify the written file still has ${MY_API_KEY}, not the resolved value
       const written = await fs.readFile(configPath, "utf-8");
@@ -72,7 +98,7 @@ describe("env snapshot TOCTOU via createConfigIO", () => {
 
       // Instance B: write WITHOUT snapshot bridging (simulates the old bug)
       const ioB = createConfigIO({ configPath, env: env as unknown as NodeJS.ProcessEnv });
-      // No setEnvSnapshot — ioB uses live env
+      // No explicit writeOptions — ioB uses live env
 
       await ioB.writeConfigFile(snapshot.config);
 
@@ -89,16 +115,15 @@ describe("env snapshot TOCTOU via createConfigIO", () => {
 
 describe("env snapshot TOCTOU via wrapper APIs", () => {
   it("uses explicit read context even if another read interleaves", async () => {
-    const prevConfigPath = process.env.OPENCLAW_CONFIG_PATH;
-    const prevCacheDisabled = process.env.OPENCLAW_DISABLE_CONFIG_CACHE;
-    const prevToken = process.env.MY_API_KEY;
     const configJson = JSON.stringify({ gateway: { remote: { token: "${MY_API_KEY}" } } }, null, 2);
-
-    try {
     await withTempConfig(configJson, async (configPath) => {
-        process.env.OPENCLAW_CONFIG_PATH = configPath;
+      await withEnvOverrides(
-        process.env.OPENCLAW_DISABLE_CONFIG_CACHE = "1";
+        {
-        process.env.MY_API_KEY = "original-key-123";
+          OPENCLAW_CONFIG_PATH: configPath,
+          OPENCLAW_DISABLE_CONFIG_CACHE: "1",
+          MY_API_KEY: "original-key-123",
+        },
+        async () => {
           const firstRead = await readConfigFileSnapshotForWrite();
           expect(firstRead.snapshot.config.gateway?.remote?.token).toBe("original-key-123");
 
@@ -112,37 +137,21 @@ describe("env snapshot TOCTOU via wrapper APIs", () => {
           const written = await fs.readFile(configPath, "utf-8");
           const parsed = JSON.parse(written);
           expect(parsed.gateway.remote.token).toBe("${MY_API_KEY}");
+        },
+      );
     });
-    } finally {
-      if (prevConfigPath === undefined) {
-        delete process.env.OPENCLAW_CONFIG_PATH;
-      } else {
-        process.env.OPENCLAW_CONFIG_PATH = prevConfigPath;
-      }
-      if (prevCacheDisabled === undefined) {
-        delete process.env.OPENCLAW_DISABLE_CONFIG_CACHE;
-      } else {
-        process.env.OPENCLAW_DISABLE_CONFIG_CACHE = prevCacheDisabled;
-      }
-      if (prevToken === undefined) {
-        delete process.env.MY_API_KEY;
-      } else {
-        process.env.MY_API_KEY = prevToken;
-      }
-    }
   });
 
   it("ignores read context when expected config path does not match", async () => {
-    const prevConfigPath = process.env.OPENCLAW_CONFIG_PATH;
-    const prevCacheDisabled = process.env.OPENCLAW_DISABLE_CONFIG_CACHE;
-    const prevToken = process.env.MY_API_KEY;
     const configJson = JSON.stringify({ gateway: { remote: { token: "${MY_API_KEY}" } } }, null, 2);
-
-    try {
     await withTempConfig(configJson, async (configPath) => {
-        process.env.OPENCLAW_CONFIG_PATH = configPath;
+      await withEnvOverrides(
-        process.env.OPENCLAW_DISABLE_CONFIG_CACHE = "1";
+        {
-        process.env.MY_API_KEY = "original-key-123";
+          OPENCLAW_CONFIG_PATH: configPath,
+          OPENCLAW_DISABLE_CONFIG_CACHE: "1",
+          MY_API_KEY: "original-key-123",
+        },
+        async () => {
           const firstRead = await readConfigFileSnapshotForWrite();
           expect(firstRead.snapshot.config.gateway?.remote?.token).toBe("original-key-123");
           expect(firstRead.writeOptions.expectedConfigPath).toBe(configPath);
@@ -156,23 +165,8 @@ describe("env snapshot TOCTOU via wrapper APIs", () => {
           const written = await fs.readFile(configPath, "utf-8");
           const parsed = JSON.parse(written);
           expect(parsed.gateway.remote.token).toBe("original-key-123");
-      });
+        },
-    } finally {
+      );
-      if (prevConfigPath === undefined) {
+    });
-        delete process.env.OPENCLAW_CONFIG_PATH;
-      } else {
-        process.env.OPENCLAW_CONFIG_PATH = prevConfigPath;
-      }
-      if (prevCacheDisabled === undefined) {
-        delete process.env.OPENCLAW_DISABLE_CONFIG_CACHE;
-      } else {
-        process.env.OPENCLAW_DISABLE_CONFIG_CACHE = prevCacheDisabled;
-      }
-      if (prevToken === undefined) {
-        delete process.env.MY_API_KEY;
-      } else {
-        process.env.MY_API_KEY = prevToken;
-      }
-    }
   });
 });

src/config/io.ts

@@ -408,6 +408,43 @@ export function parseConfigJson5(
   }
 }
 
+type ConfigReadResolution = {
+  resolvedConfigRaw: unknown;
+  envSnapshotForRestore: Record<string, string | undefined>;
+};
+
+function resolveConfigIncludesForRead(
+  parsed: unknown,
+  configPath: string,
+  deps: Required<ConfigIoDeps>,
+): unknown {
+  return resolveConfigIncludes(parsed, configPath, {
+    readFile: (candidate) => deps.fs.readFileSync(candidate, "utf-8"),
+    parseJson: (raw) => deps.json5.parse(raw),
+  });
+}
+
+function resolveConfigForRead(
+  resolvedIncludes: unknown,
+  env: NodeJS.ProcessEnv,
+): ConfigReadResolution {
+  // Apply config.env to process.env BEFORE substitution so ${VAR} can reference config-defined vars.
+  if (resolvedIncludes && typeof resolvedIncludes === "object" && "env" in resolvedIncludes) {
+    applyConfigEnv(resolvedIncludes as OpenClawConfig, env);
+  }
+
+  return {
+    resolvedConfigRaw: resolveConfigEnvVars(resolvedIncludes, env),
+    // Capture env snapshot after substitution for write-time ${VAR} restoration.
+    envSnapshotForRestore: { ...env } as Record<string, string | undefined>,
+  };
+}
+
+type ReadConfigFileSnapshotInternalResult = {
+  snapshot: ConfigFileSnapshot;
+  envSnapshotForRestore?: Record<string, string | undefined>;
+};
+
 export function createConfigIO(overrides: ConfigIoDeps = {}) {
   const deps = normalizeDeps(overrides);
   const requestedConfigPath = resolveConfigPathForDeps(deps);
@@ -417,11 +454,6 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
   const configPath =
     candidatePaths.find((candidate) => deps.fs.existsSync(candidate)) ?? requestedConfigPath;
 
-  // Snapshot of env vars captured after applyConfigEnv + resolveConfigEnvVars.
-  // Used by writeConfigFile to verify ${VAR} restoration against the env state
-  // that produced the resolved config, not the (possibly mutated) live env.
-  let envSnapshotForRestore: Record<string, string | undefined> | null = null;
-
   function loadConfig(): OpenClawConfig {
     try {
       maybeLoadDotEnvForConfig(deps.env);
@@ -439,27 +471,10 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
       }
       const raw = deps.fs.readFileSync(configPath, "utf-8");
       const parsed = deps.json5.parse(raw);
-
+      const { resolvedConfigRaw: resolvedConfig } = resolveConfigForRead(
-      // Resolve $include directives before validation
+        resolveConfigIncludesForRead(parsed, configPath, deps),
-      const resolved = resolveConfigIncludes(parsed, configPath, {
+        deps.env,
-        readFile: (p) => deps.fs.readFileSync(p, "utf-8"),
+      );
-        parseJson: (raw) => deps.json5.parse(raw),
-      });
-
-      // Apply config.env to process.env BEFORE substitution so ${VAR} can reference config-defined vars
-      if (resolved && typeof resolved === "object" && "env" in resolved) {
-        applyConfigEnv(resolved as OpenClawConfig, deps.env);
-      }
-
-      // Substitute ${VAR} env var references
-      const substituted = resolveConfigEnvVars(resolved, deps.env);
-
-      // Capture env snapshot after substitution for use by writeConfigFile.
-      // This ensures restoreEnvVarRefs verifies against the env that produced
-      // the resolved values, not a potentially mutated live env (TOCTOU fix).
-      envSnapshotForRestore = { ...deps.env } as Record<string, string | undefined>;
-
-      const resolvedConfig = substituted;
       warnOnConfigMiskeys(resolvedConfig, deps.logger);
       if (typeof resolvedConfig !== "object" || resolvedConfig === null) {
         return {};
@@ -539,7 +554,7 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
     }
   }
 
-  async function readConfigFileSnapshot(): Promise<ConfigFileSnapshot> {
+  async function readConfigFileSnapshotInternal(): Promise<ReadConfigFileSnapshotInternalResult> {
     maybeLoadDotEnvForConfig(deps.env);
     const exists = deps.fs.existsSync(configPath);
     if (!exists) {
@@ -555,6 +570,7 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
       );
       const legacyIssues: LegacyConfigIssue[] = [];
       return {
+        snapshot: {
           path: configPath,
           exists: false,
           raw: null,
@@ -566,6 +582,7 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
           issues: [],
           warnings: [],
           legacyIssues,
+        },
       };
     }
 
@@ -575,6 +592,7 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
       const parsedRes = parseConfigJson5(raw, deps.json5);
       if (!parsedRes.ok) {
         return {
+          snapshot: {
             path: configPath,
             exists: true,
             raw,
@@ -586,22 +604,21 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
             issues: [{ path: "", message: `JSON5 parse failed: ${parsedRes.error}` }],
             warnings: [],
             legacyIssues: [],
+          },
         };
       }
 
       // Resolve $include directives
       let resolved: unknown;
       try {
-        resolved = resolveConfigIncludes(parsedRes.parsed, configPath, {
+        resolved = resolveConfigIncludesForRead(parsedRes.parsed, configPath, deps);
-          readFile: (p) => deps.fs.readFileSync(p, "utf-8"),
-          parseJson: (raw) => deps.json5.parse(raw),
-        });
       } catch (err) {
         const message =
           err instanceof ConfigIncludeError
             ? err.message
             : `Include resolution failed: ${String(err)}`;
         return {
+          snapshot: {
             path: configPath,
             exists: true,
             raw,
@@ -613,26 +630,20 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
             issues: [{ path: "", message }],
             warnings: [],
             legacyIssues: [],
+          },
         };
       }
 
-      // Apply config.env to process.env BEFORE substitution so ${VAR} can reference config-defined vars
+      let readResolution: ConfigReadResolution;
-      if (resolved && typeof resolved === "object" && "env" in resolved) {
-        applyConfigEnv(resolved as OpenClawConfig, deps.env);
-      }
-
-      // Substitute ${VAR} env var references
-      let substituted: unknown;
       try {
-        substituted = resolveConfigEnvVars(resolved, deps.env);
+        readResolution = resolveConfigForRead(resolved, deps.env);
-        // Capture env snapshot (same as loadConfig — see TOCTOU comment above)
-        envSnapshotForRestore = { ...deps.env } as Record<string, string | undefined>;
       } catch (err) {
         const message =
           err instanceof MissingEnvVarError
             ? err.message
             : `Env var substitution failed: ${String(err)}`;
         return {
+          snapshot: {
             path: configPath,
             exists: true,
             raw,
@@ -644,15 +655,17 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
             issues: [{ path: "", message }],
             warnings: [],
             legacyIssues: [],
+          },
         };
       }
 
-      const resolvedConfigRaw = substituted;
+      const resolvedConfigRaw = readResolution.resolvedConfigRaw;
       const legacyIssues = findLegacyConfigIssues(resolvedConfigRaw);
 
       const validated = validateConfigObjectWithPlugins(resolvedConfigRaw);
       if (!validated.ok) {
         return {
+          snapshot: {
             path: configPath,
             exists: true,
             raw,
@@ -664,11 +677,13 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
             issues: validated.issues,
             warnings: validated.warnings,
             legacyIssues,
+          },
         };
       }
 
       warnIfConfigFromFuture(validated.config, deps.logger);
       return {
+        snapshot: {
           path: configPath,
           exists: true,
           raw,
@@ -681,7 +696,9 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
             applyTalkApiKey(
               applyModelDefaults(
                 applyAgentDefaults(
-                applySessionDefaults(applyLoggingDefaults(applyMessageDefaults(validated.config))),
+                  applySessionDefaults(
+                    applyLoggingDefaults(applyMessageDefaults(validated.config)),
+                  ),
                 ),
               ),
             ),
@@ -690,9 +707,12 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
           issues: [],
           warnings: validated.warnings,
           legacyIssues,
+        },
+        envSnapshotForRestore: readResolution.envSnapshotForRestore,
       };
     } catch (err) {
       return {
+        snapshot: {
           path: configPath,
           exists: true,
           raw: null,
@@ -704,22 +724,33 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
           issues: [{ path: "", message: `read failed: ${String(err)}` }],
           warnings: [],
           legacyIssues: [],
+        },
       };
     }
   }
 
-  async function writeConfigFile(cfg: OpenClawConfig) {
+  async function readConfigFileSnapshot(): Promise<ConfigFileSnapshot> {
+    const result = await readConfigFileSnapshotInternal();
+    return result.snapshot;
+  }
+
+  async function readConfigFileSnapshotForWrite(): Promise<ReadConfigFileSnapshotForWriteResult> {
+    const result = await readConfigFileSnapshotInternal();
+    return {
+      snapshot: result.snapshot,
+      writeOptions: {
+        envSnapshotForRestore: result.envSnapshotForRestore,
+        expectedConfigPath: configPath,
+      },
+    };
+  }
+
+  async function writeConfigFile(cfg: OpenClawConfig, options: ConfigWriteOptions = {}) {
     clearConfigCache();
     let persistCandidate: unknown = cfg;
-    // Save the injected env snapshot before readConfigFileSnapshot() overwrites it
+    const { snapshot } = await readConfigFileSnapshotInternal();
-    // with a new snapshot based on the (possibly mutated) live env.
-    const savedEnvSnapshot = envSnapshotForRestore;
-    const snapshot = await readConfigFileSnapshot();
     let envRefMap: Map<string, string> | null = null;
     let changedPaths: Set<string> | null = null;
-    if (savedEnvSnapshot) {
-      envSnapshotForRestore = savedEnvSnapshot;
-    }
     if (snapshot.valid && snapshot.exists) {
       const patch = createMergePatch(snapshot.config, cfg);
       persistCandidate = applyMergePatch(snapshot.resolved, patch);
@@ -771,7 +802,7 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
           // Use env snapshot from when config was loaded (if available) to avoid
           // TOCTOU issues where env changes between load and write. Falls back to
           // live env if no snapshot exists (e.g., first write before any load).
-          const envForRestore = envSnapshotForRestore ?? deps.env;
+          const envForRestore = options.envSnapshotForRestore ?? deps.env;
           cfgToWrite = restoreEnvVarRefs(
             cfgToWrite,
             parsedRes.parsed,
@@ -851,15 +882,8 @@ export function createConfigIO(overrides: ConfigIoDeps = {}) {
     configPath,
     loadConfig,
     readConfigFileSnapshot,
+    readConfigFileSnapshotForWrite,
     writeConfigFile,
-    /** Return the env snapshot captured during the last loadConfig/readConfigFileSnapshot, or null. */
-    getEnvSnapshot(): Record<string, string | undefined> | null {
-      return envSnapshotForRestore;
-    },
-    /** Inject an env snapshot (e.g. from a prior IO instance) for use by writeConfigFile. */
-    setEnvSnapshot(snapshot: Record<string, string | undefined>): void {
-      envSnapshotForRestore = snapshot;
-    },
   };
 }
 
@@ -928,27 +952,17 @@ export async function readConfigFileSnapshot(): Promise<ConfigFileSnapshot> {
 }
 
 export async function readConfigFileSnapshotForWrite(): Promise<ReadConfigFileSnapshotForWriteResult> {
-  const io = createConfigIO();
+  return await createConfigIO().readConfigFileSnapshotForWrite();
-  const snapshot = await io.readConfigFileSnapshot();
-  return {
-    snapshot,
-    writeOptions: {
-      envSnapshotForRestore: io.getEnvSnapshot() ?? undefined,
-      expectedConfigPath: io.configPath,
-    },
-  };
 }
 
 export async function writeConfigFile(
   cfg: OpenClawConfig,
   options: ConfigWriteOptions = {},
 ): Promise<void> {
-  clearConfigCache();
   const io = createConfigIO();
   const sameConfigPath =
     options.expectedConfigPath === undefined || options.expectedConfigPath === io.configPath;
-  if (sameConfigPath && options.envSnapshotForRestore) {
+  await io.writeConfigFile(cfg, {
-    io.setEnvSnapshot(options.envSnapshotForRestore);
+    envSnapshotForRestore: sameConfigPath ? options.envSnapshotForRestore : undefined,
-  }
+  });
-  await io.writeConfigFile(cfg);
 }