Merge #767 fix(custom-model): give a session its own env values back when the selection goes

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

src/custom-model-injection-apply.ts

@@ -301,6 +301,8 @@ export interface AppliedCustomModel {
   envKeys: string[];
   configDir?: string;
   launchModel?: string;
+  /** The env key among `envKeys` that moves the CLI's home, see `CustomModelBookkeeping`. */
+  homeEnvVar?: string;
 }
 
 /**
@@ -360,9 +362,16 @@ export function applyCustomModelInjection(
       envKeys: Object.keys(envOverrides),
       configDir,
       launchModel: injection.launchModel,
+      homeEnvVar: injection.configDirVar,
     };
   }
   const configDir = customModelConfigDir(sessionId);
   const envOverrides = applyConfigDirInjection(configDir, injection, sessionEnvLookup?.(injection.dirEnvVar));
-  return { envOverrides, envKeys: Object.keys(envOverrides), configDir, launchModel: injection.launchModel };
+  return {
+    envOverrides,
+    envKeys: Object.keys(envOverrides),
+    configDir,
+    launchModel: injection.launchModel,
+    homeEnvVar: injection.dirEnvVar,
+  };
 }

src/session-env-clamp.ts

@@ -26,13 +26,13 @@
  * re-resolves the grant at the moment it builds a session, for the same reason.
  *
  * @dependencies user-store (canUsernameRunPrivilegedCommands), config/cli-registry
- * @consumedby web/routes/session-routes, web/routes/reboot-restore-routes
+ * @consumedby web/routes/session-routes, web/routes/reboot-restore-routes, web/server
  *
  * @module session-env-clamp
  */
 
 import { canUsernameRunPrivilegedCommands } from './user-store.js';
-import { enabledClis } from './config/cli-registry/registry.js';
+import { listClis } from './config/cli-registry/registry.js';
 
 /**
  * Env-var keys a non-granted owner must not be able to set, because each one
@@ -68,7 +68,8 @@ import { enabledClis } from './config/cli-registry/registry.js';
  *   already allowlisted for pi and not addressed here — see resolveOmpHome()).
  */
 export function ownerClampedEnvKeys(): string[] {
-  return enabledClis().flatMap((entry) => entry.capabilities.privilegedEnvKeys);
+  // Every CLI, enabled or not: a session can outlive the switch that disabled its CLI.
+  return listClis().flatMap((entry) => entry.capabilities.privilegedEnvKeys);
 }
 
 /**

src/session.ts

@@ -29,7 +29,7 @@
  */
 
 import { EventEmitter } from 'node:events';
-import { basename } from 'node:path';
+import { basename, isAbsolute } from 'node:path';
 import { realpath } from 'node:fs/promises';
 import { execSync, execFileSync } from 'node:child_process';
 import { v4 as uuidv4 } from 'uuid';
@@ -1372,6 +1372,17 @@ export class Session extends EventEmitter {
     return value && value.trim() ? value.trim() : undefined;
   }
 
+  /**
+   * What the SESSION itself set for `key`, same single-key rule. `envOverrideValue()` answers
+   * with whatever the map holds, which under a selection is the injected value: a re-apply
+   * resolving the CLI's real home must read the session's own home, or none.
+   */
+  ownEnvValue(key: string): string | undefined {
+    const selection = this._customModel;
+    const value = selection?.envKeys.includes(key) ? selection.ownEnv?.[key] : this._envOverrides?.[key];
+    return value?.trim() || undefined;
+  }
+
   /** Owning username in multi-user mode, else undefined. */
   get owner(): string | undefined {
     return this._owner;
@@ -1992,21 +2003,38 @@ export class Session extends EventEmitter {
 
   /**
    * The full selection incl. bookkeeping, for state.json ONLY (`__customModel`, the
-   * same disk-only convention as `getEnvOverridesForPersist()`). Carries no env values,
-   * so nothing secret lands on disk; recovery re-derives them from the endpoint store.
+   * same disk-only convention as `getEnvOverridesForPersist()`). Carries no injected env
+   * values, so nothing secret lands on disk; recovery re-derives them from the endpoint store.
    * Without this a Codeman restart left the pane on the custom endpoint (tmux keeps
    * its `setenv`s) while `customModel` came back undefined, so the state was wrong and
    * clearing had nothing to unset. Must NOT be included in any API-bound serializer.
+   *
+   * The one value it does carry is the session's own home directory (`ownEnv`, narrowed to
+   * `homeEnvVar`), a path like `configDir` beside it: the selection overwrote the tmux
+   * `setenv` that held it, so without this a clear after a restart would lose it. Any other
+   * entry of `ownEnv` may be a credential and stays in memory, and so does a `homeEnvVar`
+   * value that is not an absolute path: the registry types that name as any env var, so a
+   * `clis.json` override can point it at a credential.
    */
   getCustomModelForPersist(): CustomModelBookkeeping | undefined {
-    return this._customModel ? { ...this._customModel, envKeys: [...this._customModel.envKeys] } : undefined;
+    if (!this._customModel) return undefined;
+    const { ownEnv, ...selection } = this._customModel;
+    const { homeEnvVar } = selection;
+    const ownHome = homeEnvVar ? ownEnv?.[homeEnvVar] : undefined;
+    return {
+      ...selection,
+      envKeys: [...selection.envKeys],
+      ...(homeEnvVar && ownHome && isAbsolute(ownHome) ? { ownEnv: { [homeEnvVar]: ownHome } } : {}),
+    };
   }
 
   /**
    * Update this session's custom-model selection and merge the endpoint's injected env
    * vars into `_envOverrides` — first UNDOING whatever the previous selection injected
-   * (removing exactly those env keys), so switching endpoints, or clearing back to the
-   * harness's native cloud default, never leaves a stale key behind. Synchronous and
+   * (removing exactly those env keys, and putting back a value the session had set under
+   * one itself), so switching endpoints, or clearing back to the harness's native cloud
+   * default, never leaves a stale key behind and never loses the session's own `DSH_HOME`
+   * or `CLAUDE_CONFIG_DIR`. Synchronous and
    * side-effect-free beyond mutating state, matching `setNice`/`setColor` above — this
    * class does no file IO, so it reports the PREVIOUS `configDir` (if any) for the
    * caller to clean up on disk (custom-model-injection.ts's configDir kind).
@@ -2016,7 +2044,12 @@ export class Session extends EventEmitter {
    * `_buildRespawnPaneOptions().unsetEnvKeys`): the tmux session inherits every
    * `setenv` into `respawn-pane`, so dropping them from the map alone would relaunch
    * the CLI still pointed at the old endpoint — and for the `configDir` kinds, at a
-   * `HOME`/`CODEX_HOME`/`GROK_HOME` the caller has just deleted.
+   * `HOME`/`CODEX_HOME`/`GROK_HOME` the caller has just deleted. A key put back stays out
+   * of that queue: `applyEnvOverrides()` writes its value over the injected one.
+   *
+   * `next.ownEnv` is for recovery, which states the session's own values itself: after a
+   * restart `_envOverrides` holds what the previous run injected (`CLAUDE_CONFIG_DIR` is
+   * persisted), and reading those as the session's own would pin it to a deleted directory.
    */
   setCustomModel(
     next: CustomModelBookkeeping | undefined,
@@ -2025,13 +2058,30 @@ export class Session extends EventEmitter {
     const previousConfigDir = this._customModel?.configDir;
     const removedEnvKeys: string[] = [];
     if (this._customModel) {
+      const previousOwnEnv = this._customModel.ownEnv ?? {};
       for (const key of this._customModel.envKeys) {
         if (this._envOverrides) delete this._envOverrides[key];
         removedEnvKeys.push(key);
-        this._pendingEnvUnsets.add(key);
+        const own = previousOwnEnv[key];
+        if (own) {
+          this._envOverrides = { ...(this._envOverrides ?? {}), [key]: own };
+          this._pendingEnvUnsets.delete(key);
+        } else {
+          this._pendingEnvUnsets.add(key);
+        }
       }
     }
-    this._customModel = next ? { ...next, envKeys: [...next.envKeys] } : undefined;
+    // Read after the undo, so a switch between endpoints carries the own values across. An
+    // empty value is skipped: `applyEnvOverrides()` never set one, so it has nothing to
+    // restore, and the injected key still has to be unset.
+    const ownEnv: Record<string, string> = { ...next?.ownEnv };
+    if (!next?.ownEnv) {
+      for (const key of Object.keys(envOverrides ?? {})) {
+        const own = this._envOverrides?.[key];
+        if (own) ownEnv[key] = own;
+      }
+    }
+    this._customModel = next ? { ...next, envKeys: [...next.envKeys], ownEnv } : undefined;
     if (envOverrides && Object.keys(envOverrides).length > 0) {
       this._envOverrides = { ...(this._envOverrides ?? {}), ...envOverrides };
       // A key the new selection sets again does not need an unset (applyEnvOverrides

src/types/session.ts

@@ -702,6 +702,20 @@ export interface CustomModelBookkeeping extends CustomModelSelection {
   configDir?: string;
   /** Value forced onto the CLI's `model` launch param (pi/omp `custom/<id>`, grok's block name). */
   launchModel?: string;
+  /**
+   * The variable that moves the CLI's home, when the recipe redirects one (`DSH_HOME`,
+   * `CLAUDE_CONFIG_DIR`, `CODEX_HOME`, `GROK_HOME`). Its value is a path, which is what
+   * lets {@link CustomModelBookkeeping.ownEnv} keep that one entry on disk.
+   */
+  homeEnvVar?: string;
+  /**
+   * What the session itself had set under keys the selection replaced, put back when the
+   * selection goes. In memory this can hold a credential (a session's own `GEMINI_API_KEY`),
+   * so `Session.getCustomModelForPersist()` writes `homeEnvVar`'s entry, when it is an absolute
+   * path, and nothing else: tmux keeps only the injected value, so state.json is the one copy
+   * left after a restart.
+   */
+  ownEnv?: Record<string, string>;
 }
 
 /**

src/web/routes/session-routes.ts

@@ -202,7 +202,7 @@ import {
 } from '../../docker-hosts.js';
 import { resolveCaseRepository, resolveManagedWorktree } from '../../git-worktrees.js';
 import { LRUMap } from '../../utils/lru-map.js';
-import { describeUnknownPath, probePathKind } from '../../utils/index.js';
+import { describeUnknownPath, probePathKind, type PathProbeOptions } from '../../utils/index.js';
 import { findLatestOmpSessionId } from '../../utils/omp-session-resolver.js';
 import { scanOmpSessionsHistory } from '../../omp-transcript.js';
 import { scanCodexSessionsHistory, codexThreadBySessionId } from '../../codex-transcript.js';
@@ -244,18 +244,20 @@ function landedInGroup(layout: TabLayout | undefined, sessionId: string): boolea
  * Probe the session's own home with `probePathKind(..., { pastCap: true })` first: on
  * `'unknown'` (a dead mount that never answered) hand the apply a lookup that hides that
  * one env var, so linking falls through to the server's own environment and the unresponsive
- * path is never touched. Any other entry's lookup passes through untouched.
+ * path is never touched. Any other entry's lookup passes through untouched. Boot recovery
+ * passes `probeOptions` of `{}`: it walks every saved session, so it keeps the bulk cap.
  */
 export async function guardSessionConfigDirHome(
   entry: Pick<CliEntry, 'capabilities'>,
-  baseLookup: (envVar: string) => string | undefined
+  baseLookup: (envVar: string) => string | undefined,
+  probeOptions: PathProbeOptions = { pastCap: true }
 ): Promise<(envVar: string) => string | undefined> {
   const recipe = entry.capabilities.customModelInjection;
   if (recipe?.kind !== 'configDir' || !recipe.linkFromRealHome?.length) return baseLookup;
   const dirEnvVar = recipe.dirEnvVar;
   const supplied = baseLookup(dirEnvVar)?.trim();
   if (!supplied || !isAbsolute(supplied)) return baseLookup;
-  if ((await probePathKind(supplied, { pastCap: true })) === 'unknown') {
+  if ((await probePathKind(supplied, probeOptions)) === 'unknown') {
     console.warn(
       `[custom-model] ${dirEnvVar} (${supplied}) did not answer; linking from the server's own home instead`
     );
@@ -1675,7 +1677,7 @@ export function registerSessionRoutes(
       session.id,
       contextLength,
       session.workingDir,
-      await guardSessionConfigDirHome(entry, (envVar) => session.envOverrideValue(envVar))
+      await guardSessionConfigDirHome(entry, (envVar) => session.ownEnvValue(envVar))
     );
     if (!applied) {
       return createErrorResponse(ApiErrorCode.OPERATION_FAILED, `${session.mode} has no known custom-model mechanism`);
@@ -1689,6 +1691,7 @@ export function registerSessionRoutes(
         envKeys: applied.envKeys,
         configDir: applied.configDir,
         launchModel: applied.launchModel,
+        homeEnvVar: applied.homeEnvVar,
       },
       applied.envOverrides
     );
@@ -4467,12 +4470,12 @@ export function registerSessionRoutes(
     // unknown endpoint, a model id the CLI's argv pattern can't carry) rather than trusting
     // a lighter version of them, since this is the same server-side authority reached a
     // different way, not a separate, less-checked path.
-    let qsCustomModelEnvOverrides = qsGatedEnvOverrides;
-    // Only the INJECTED keys (never the caller's envOverrides merged in) — this is what
-    // setCustomModel() bookkeeping must be given below. The Session constructor already
-    // applies qsCustomModelEnvOverrides (the full merged set) directly; re-merging that
-    // full set into setCustomModel() would put CLAUDE_CODE_EFFORT_LEVEL back after the
-    // constructor stripped it (see setCustomModel()'s own doc comment in session.ts).
+    //
+    // Only the INJECTED keys (never the caller's envOverrides merged in): setCustomModel()
+    // merges these over what the Session constructor took, which is how it learns which of
+    // the caller's values the selection replaced. Handing the constructor the merged set
+    // would hide them, and giving setCustomModel() that set would put
+    // CLAUDE_CODE_EFFORT_LEVEL back after the constructor stripped it.
     let qsCustomModelAppliedEnvOverrides: Record<string, string> | undefined;
     let qsCustomModelLaunchModel: string | undefined;
     let qsCustomModelSessionId: string | undefined;
@@ -4485,6 +4488,7 @@ export function registerSessionRoutes(
           envKeys: string[];
           configDir?: string;
           launchModel?: string;
+          homeEnvVar?: string;
         }
       | undefined;
     if (customModel) {
@@ -4588,7 +4592,6 @@ export function registerSessionRoutes(
       if (!cmApplied) {
         return createErrorResponse(ApiErrorCode.OPERATION_FAILED, `${mode} has no known custom-model mechanism`);
       }
-      qsCustomModelEnvOverrides = { ...qsGatedEnvOverrides, ...cmApplied.envOverrides };
       qsCustomModelAppliedEnvOverrides = cmApplied.envOverrides;
       qsCustomModelLaunchModel = cmApplied.launchModel;
       qsCustomModelBookkeeping = {
@@ -4598,6 +4601,7 @@ export function registerSessionRoutes(
         envKeys: cmApplied.envKeys,
         configDir: cmApplied.configDir,
         launchModel: cmApplied.launchModel,
+        homeEnvVar: cmApplied.homeEnvVar,
       };
 
       // Actually kick off llama-swap's load now — see the dedicated apply route's own
@@ -4630,7 +4634,7 @@ export function registerSessionRoutes(
       copilotConfig: mode === 'copilot' ? qsGatedCopilotConfig : undefined,
       deepSeekConfig: mode === 'deepseek' ? qsGatedDeepSeekConfig : undefined,
       ompConfig: qsResolvedOmpConfig,
-      envOverrides: qsCustomModelEnvOverrides,
+      envOverrides: qsGatedEnvOverrides,
       effort,
       advisorModel,
       remote,
@@ -4659,10 +4663,10 @@ export function registerSessionRoutes(
     const session = new Session(qsSessionOptions);
 
     // Records the selection for session.customModel/getCustomModelForPersist() and future
-    // clear/switch calls — the actual env vars and launch-model config are already part of
-    // the launch above (constructor envOverrides, piConfig/grokConfig/ompConfig.model), so
-    // this is bookkeeping only, never a restart: setCustomModel() is synchronous state, no
-    // tmux IO of its own (see its own doc comment in session.ts).
+    // clear/switch calls, and merges the injected env, which the launch below reads along
+    // with the launch-model config set above (piConfig/grokConfig/ompConfig.model). Nothing
+    // has started yet, so this is never a restart: setCustomModel() is synchronous state,
+    // no tmux IO of its own (see its own doc comment in session.ts).
     if (qsCustomModelBookkeeping) {
       session.setCustomModel(qsCustomModelBookkeeping, qsCustomModelAppliedEnvOverrides);
     }

src/web/server.ts

@@ -104,6 +104,7 @@ import { RunSummaryTracker } from '../run-summary.js';
 import { PlanOrchestrator } from '../plan-orchestrator.js';
 import { OrchestratorLoop } from '../orchestrator-loop.js';
 import { getLifecycleLog } from '../session-lifecycle-log.js';
+import { clampEnvOverridesForOwner } from '../session-env-clamp.js';
 import {
   applyWorkspaceHooks,
   pruneAgentSessionPreambles,
@@ -135,6 +136,7 @@ import { fetchWorkspaceOnce } from '../git-fetch.js';
 import { GitFetchScheduler } from '../git-fetch-scheduler.js';
 import { probePathKind } from '../utils/bounded-path-probe.js';
 import { defaultDockerWorkspaces } from './routes/git-status-routes.js';
+import { guardSessionConfigDirHome } from './routes/session-routes.js';
 import { composerIsEmpty, plainFrame } from '../session-submit-verifier.js';
 import { detectUsageLimitPause } from '../usage-limit-patterns.js';
 import { EditHistorySync } from './edit-history-sync.js';
@@ -1519,7 +1521,8 @@ export class WebServer extends EventEmitter {
    */
   private async _rebuildCustomModelEnv(
     session: Session,
-    saved: CustomModelBookkeeping
+    saved: CustomModelBookkeeping,
+    ownEnv: Record<string, string> = {}
   ): Promise<AppliedCustomModel | undefined> {
     const entry = getCli(session.mode);
     if (!entry) return undefined;
@@ -1531,19 +1534,21 @@ export class WebServer extends EventEmitter {
       return undefined;
     }
     try {
-      // Boot recovery omits both `workingDir` (no reason to re-answer a trust dialog that
-      // already fired once) and the session-env real-home lookup: the injected VALUES are
-      // never persisted, so after a real restart the session can no longer carry its
-      // allowlisted `DSH_HOME` either (`getEnvOverridesForPersist()` keeps only
-      // `CLAUDE_CODE_*`/`CLAUDE_CONFIG_DIR`): the lookup could only ever hit on the first
-      // apply in-process. Real-home resolution therefore falls back to the SERVER's
+      // Boot recovery omits `workingDir` (no reason to re-answer a trust dialog that already
+      // fired once). The real-home lookup reads the record's own values, since the session's
+      // overrides hold what the previous run injected. `linkRealHomeEntry()` keeps what is
+      // there, so this only decides an entry that is missing, which an own home skips when it
+      // does not hold one. With no own home, or one that does not answer (probed under the
+      // bulk cap, this walks every saved session), a missing entry links from the SERVER's
       // `DSH_HOME` or the recipe's `realHomeFallback` (~/.dsh).
       return applyCustomModelInjection(
         entry,
         endpoint,
         saved.modelId,
         session.id,
-        resolveCustomModelContextLength(endpoint, saved.modelId)
+        resolveCustomModelContextLength(endpoint, saved.modelId),
+        undefined,
+        await guardSessionConfigDirHome(entry, (envVar) => ownEnv[envVar], {})
       );
     } catch (err) {
       console.warn('[WebServer] Failed to rebuild custom-model env on recovery:', err);
@@ -1559,18 +1564,25 @@ export class WebServer extends EventEmitter {
    * while the new recipe writes one, and clearing it later could then never delete the
    * generated dir. `envKeys` is the UNION of saved and applied keys: the tmux session may
    * still carry `setenv`s from both, and a clear has to unset every one.
+   *
+   * `ownEnv` is always passed, empty for a record saved before it existed: `_envOverrides`
+   * holds what the previous run injected, and only this record says what was the session's
+   * own. It meets the owner's grant again, like every other env override rebuilt from disk.
    */
   private async _restoreCustomModel(session: Session, saved: CustomModelBookkeeping): Promise<void> {
-    const applied = await this._rebuildCustomModelEnv(session, saved);
+    const ownEnv = (await clampEnvOverridesForOwner(session.owner, saved.ownEnv)) ?? {};
+    const applied = await this._rebuildCustomModelEnv(session, saved, ownEnv);
     if (!applied) {
       // Endpoint gone or CLI unregistered: restore the bookkeeping exactly as saved, which
       // is what a later clear needs.
-      session.setCustomModel(saved);
+      session.setCustomModel({ ...saved, ownEnv });
       return;
     }
     session.setCustomModel(
       {
         ...saved,
+        ownEnv,
+        homeEnvVar: applied.homeEnvVar ?? saved.homeEnvVar,
         // Prefer the current recipe's dir; keep the saved one when the recipe writes none,
         // so a clear can still delete a directory an older recipe created.
         configDir: applied.configDir ?? saved.configDir,
@@ -1590,8 +1602,9 @@ export class WebServer extends EventEmitter {
     // __attachmentHistory keeps the private (externalPath-bearing) history on disk,
     // separate from the sanitized public attachmentHistory in toState().
     const attachmentHistory = session.getAttachmentHistoryForPersist();
-    // __customModel keeps the selection's bookkeeping (injected env KEYS, config dir,
-    // launch model; never the values) so recovery can restore and later clear it.
+    // __customModel keeps the selection's bookkeeping (injected env KEYS, config dir, launch
+    // model, the session's own home path; never an injected value or any other own value)
+    // so recovery can restore and later clear it.
     const customModel = session.getCustomModelForPersist();
     const state = {
       ...base,