test: add snapshotEnv helper and restore leaked env and auth state (TA-001 slice 4) (#2335)

CONTRIBUTING.md

@@ -160,6 +160,7 @@ Dependency changes need a clear project benefit — fixing a bug, addressing a s
 - Preserve existing repo patterns unless the change is intentionally refactoring them.
 - Add or update tests when the change affects behavior.
 - For vendor/integration catalog tests, cover each vendor with one table-driven test of its registry contract (route resolution, default base URL and model, credential env vars, catalog-to-descriptor linkage, and capability/limit values the runtime uses), plus a small number of behavioral tests such as model-discovery keep/drop parsing or availability windows. Do not copy full curated model lists or cache/refresh settings from vendor definitions into tests.
+- When tests set or delete environment variables, import `snapshotEnv` from `src/test/envSnapshot.ts` using the appropriate relative `.js` path. In `beforeEach`, save the named variables before changing them with `restoreEnv = snapshotEnv(keys)`; pass `{ clear: true }` when setup needs those variables unset. Call `restoreEnv()` in `afterEach`, or use `finally` for changes scoped to one test. Restoration preserves previous values, including empty strings, and deletes variables that were previously unset. List only the variables the test changes or deliberately clears; do not replace `process.env` or blanket-clear unrelated variables. When using the shared mutation lock, snapshot after acquiring it and restore before releasing it.
 - In timing-sensitive tests, drive timers with Bun's fake timers (`jest.useFakeTimers()` and `jest.advanceTimersByTime()`) instead of racing real delays, and restore real timers during cleanup, so the tests pass in a single-process whole-suite run as well as in the sharded runner.
 - Await completion promises or explicit readiness signals before asserting asynchronous results. For tests that deliberately suppress readiness, use a short, bounded wait and assert the expected failure.
 - For tests that launch real subprocesses, set an explicit per-test timeout longer than the child-process deadline, leaving time for assertions and cleanup. The launcher heap tests use an 8-second child deadline and a 10-second test timeout. Document the reason for the timeout beside the test.

src/hooks/useApiKeyVerification.test.tsx

@@ -1,6 +1,6 @@
 import { PassThrough } from 'node:stream'
 
-import { afterEach, beforeEach, expect, mock, test } from 'bun:test'
+import { afterAll, afterEach, beforeEach, expect, mock, test } from 'bun:test'
 import React from 'react'
 import { createRoot, Text } from '../ink.js'
 import * as realState from '../bootstrap/state.js'
@@ -9,7 +9,11 @@ import {
   acquireSharedMutationLock,
   releaseSharedMutationLock,
 } from '../test/sharedMutationLock.js'
-import * as realAuth from '../utils/auth.js'
+
+const nonce = `${Date.now()}-${Math.random()}`
+const realAuth = {
+  ...await import(`../utils/auth.js?real=${nonce}`),
+}
 
 type AuthState = {
   anthropicAuthEnabled: boolean
@@ -75,6 +79,10 @@ afterEach(() => {
   }
 })
 
+afterAll(() => {
+  mock.module('../utils/auth.js', () => realAuth)
+})
+
 test('useApiKeyVerification resets stale missing status when the session switches to a third-party provider', async () => {
   const authState: AuthState = {
     anthropicAuthEnabled: true,

src/services/api/openaiShim.ollamaTextToolCalls.test.ts

@@ -10,9 +10,12 @@
  *   5. P1 context guard — bare JSON in explanatory prose is skipped
  */
 import { afterEach, beforeEach, describe, expect, test } from 'bun:test'
+import { snapshotEnv } from '../../test/envSnapshot.js'
 import { createOpenAIShimClient } from './openaiShim.js'
 import { parseTextToolCalls } from './openaiShim/rawToolCallParsing.js'
 
+const ENV_KEYS = ['OLLAMA_BASE_URL', 'OPENAI_API_KEY', 'OPENAI_BASE_URL'] as const
+
 type FetchType = typeof globalThis.fetch
 
 type OpenAIShimClient = {
@@ -224,17 +227,17 @@ const ollamaToolChunk = (toolCalls: unknown[], finishReason?: string) => ({
 
 describe('Ollama streaming — think-tag filtering on text-tool fallback (P1)', () => {
   let originalFetch: FetchType
+  let restoreEnv: () => void
   beforeEach(() => {
     originalFetch = globalThis.fetch
+    restoreEnv = snapshotEnv(ENV_KEYS)
     process.env.OLLAMA_BASE_URL = 'http://localhost:11434'
     process.env.OPENAI_API_KEY = 'test-key'
     process.env.OPENAI_BASE_URL = 'http://localhost:11434/v1'
   })
   afterEach(() => {
     globalThis.fetch = originalFetch
-    delete process.env.OLLAMA_BASE_URL
-    delete process.env.OPENAI_API_KEY
-    delete process.env.OPENAI_BASE_URL
+    restoreEnv()
   })
 
   test('<think> content is NOT emitted as assistant text when text-tool fallback fires', async () => {
@@ -278,17 +281,17 @@ describe('Ollama streaming — think-tag filtering on text-tool fallback (P1)',
 
 describe('Ollama streaming — plain text response with no tool calls', () => {
   let originalFetch: FetchType
+  let restoreEnv: () => void
   beforeEach(() => {
     originalFetch = globalThis.fetch
+    restoreEnv = snapshotEnv(ENV_KEYS)
     process.env.OLLAMA_BASE_URL = 'http://localhost:11434'
     process.env.OPENAI_API_KEY = 'test-key'
     process.env.OPENAI_BASE_URL = 'http://localhost:11434/v1'
   })
   afterEach(() => {
     globalThis.fetch = originalFetch
-    delete process.env.OLLAMA_BASE_URL
-    delete process.env.OPENAI_API_KEY
-    delete process.env.OPENAI_BASE_URL
+    restoreEnv()
   })
 
   test('plain text in two chunks (content then stop) is emitted as text_delta', async () => {
@@ -385,17 +388,17 @@ describe('Ollama streaming — plain text response with no tool calls', () => {
 
 describe('Ollama streaming — visible text before real structured tool_calls (P2)', () => {
   let originalFetch: FetchType
+  let restoreEnv: () => void
   beforeEach(() => {
     originalFetch = globalThis.fetch
+    restoreEnv = snapshotEnv(ENV_KEYS)
     process.env.OLLAMA_BASE_URL = 'http://localhost:11434'
     process.env.OPENAI_API_KEY = 'test-key'
     process.env.OPENAI_BASE_URL = 'http://localhost:11434/v1'
   })
   afterEach(() => {
     globalThis.fetch = originalFetch
-    delete process.env.OLLAMA_BASE_URL
-    delete process.env.OPENAI_API_KEY
-    delete process.env.OPENAI_BASE_URL
+    restoreEnv()
   })
 
   test('visible assistant text is preserved when real delta.tool_calls follow it', async () => {
@@ -441,17 +444,17 @@ describe('Ollama streaming — visible text before real structured tool_calls (P
 
 describe('Ollama streaming — visible prose before text-based tool-call fallback (P2 buffered path)', () => {
   let originalFetch: FetchType
+  let restoreEnv: () => void
   beforeEach(() => {
     originalFetch = globalThis.fetch
+    restoreEnv = snapshotEnv(ENV_KEYS)
     process.env.OLLAMA_BASE_URL = 'http://localhost:11434'
     process.env.OPENAI_API_KEY = 'test-key'
     process.env.OPENAI_BASE_URL = 'http://localhost:11434/v1'
   })
   afterEach(() => {
     globalThis.fetch = originalFetch
-    delete process.env.OLLAMA_BASE_URL
-    delete process.env.OPENAI_API_KEY
-    delete process.env.OPENAI_BASE_URL
+    restoreEnv()
   })
 
   test('visible prose before text-based JSON tool call is preserved in emitted text_delta', async () => {
@@ -513,17 +516,17 @@ describe('parseTextToolCalls — pretty-printed bare JSON detection', () => {
 
 describe('Ollama streaming — non-stop terminal finish reasons flush buffer', () => {
   let originalFetch: FetchType
+  let restoreEnv: () => void
   beforeEach(() => {
     originalFetch = globalThis.fetch
+    restoreEnv = snapshotEnv(ENV_KEYS)
     process.env.OLLAMA_BASE_URL = 'http://localhost:11434'
     process.env.OPENAI_API_KEY = 'test-key'
     process.env.OPENAI_BASE_URL = 'http://localhost:11434/v1'
   })
   afterEach(() => {
     globalThis.fetch = originalFetch
-    delete process.env.OLLAMA_BASE_URL
-    delete process.env.OPENAI_API_KEY
-    delete process.env.OPENAI_BASE_URL
+    restoreEnv()
   })
 
   test('buffered text is flushed when finish_reason is "length"', async () => {

src/services/api/openaiShim/xmlToolCallParsing.test.ts

@@ -15,6 +15,7 @@
  *   D. <tool_calls:ID><tool_call:ID>NAME<parameter name="KEY">VALUE</parameter>…</tool_calls:ID>
  */
 import { afterEach, beforeEach, describe, expect, test } from 'bun:test'
+import { snapshotEnv } from '../../../test/envSnapshot.js'
 import { createOpenAIShimClient } from '../openaiShim.js'
 import {
   parseXmlToolCalls as parseXmlToolCallsModule,
@@ -283,16 +284,17 @@ describe('parseXmlToolCalls', () => {
 
 describe('GLM streaming — XML tool calls', () => {
   let originalFetch: FetchType
+  let restoreEnv: () => void
   beforeEach(() => {
     originalFetch = globalThis.fetch
+    restoreEnv = snapshotEnv(['OPENAI_API_KEY', 'OPENAI_BASE_URL'])
     process.env.OPENAI_API_KEY = 'test-key'
     // Deliberately NOT an Ollama endpoint, so isOllama=false.
     process.env.OPENAI_BASE_URL = 'https://gateway.example.com/v1'
   })
   afterEach(() => {
     globalThis.fetch = originalFetch
-    delete process.env.OPENAI_API_KEY
-    delete process.env.OPENAI_BASE_URL
+    restoreEnv()
   })
 
   async function run(

src/test/envSnapshot.test.ts

@@ -0,0 +1,97 @@
+import { describe, expect, test } from 'bun:test'
+import { snapshotEnv } from './envSnapshot.js'
+
+const KEYS = [
+  'TA001_ENV_SNAPSHOT_TEST_POPULATED',
+  'TA001_ENV_SNAPSHOT_TEST_EMPTY',
+  'TA001_ENV_SNAPSHOT_TEST_UNDEFINED',
+  'TA001_ENV_SNAPSHOT_TEST_ABSENT',
+  'TA001_ENV_SNAPSHOT_TEST_UNLISTED',
+] as const
+
+function withTestEnv(run: () => void): void {
+  const incoming = KEYS.map(key => [key, process.env[key]] as const)
+  try {
+    run()
+  } finally {
+    for (const [key, value] of incoming) {
+      if (value === undefined) {
+        delete process.env[key]
+      } else {
+        process.env[key] = value
+      }
+    }
+  }
+}
+
+describe('snapshotEnv', () => {
+  test('leaves values unchanged by default and restores exact strings', () => {
+    withTestEnv(() => {
+      const [populated, empty, literalUndefined] = KEYS
+      process.env[populated] = 'original'
+      process.env[empty] = ''
+      process.env[literalUndefined] = 'undefined'
+
+      const restoreEnv = snapshotEnv([populated, empty, literalUndefined])
+      expect(process.env[populated]).toBe('original')
+      expect(process.env[empty]).toBe('')
+      expect(process.env[literalUndefined]).toBe('undefined')
+
+      process.env[populated] = 'changed'
+      process.env[empty] = 'changed'
+      delete process.env[literalUndefined]
+      restoreEnv()
+
+      expect(process.env[populated]).toBe('original')
+      expect(process.env[empty]).toBe('')
+      expect(process.env[literalUndefined]).toBe('undefined')
+    })
+  })
+
+  test('clears only named keys and saves their values before clearing', () => {
+    withTestEnv(() => {
+      const [populated, empty, , , unlisted] = KEYS
+      process.env[populated] = 'original'
+      process.env[empty] = ''
+      process.env[unlisted] = 'untouched'
+
+      const restoreEnv = snapshotEnv([populated, empty], { clear: true })
+      expect(Object.hasOwn(process.env, populated)).toBe(false)
+      expect(Object.hasOwn(process.env, empty)).toBe(false)
+      expect(process.env[unlisted]).toBe('untouched')
+
+      restoreEnv()
+      expect(process.env[populated]).toBe('original')
+      expect(process.env[empty]).toBe('')
+      expect(process.env[unlisted]).toBe('untouched')
+    })
+  })
+
+  test('deletes keys that were previously absent', () => {
+    withTestEnv(() => {
+      const absent = KEYS[3]
+      delete process.env[absent]
+      const restoreEnv = snapshotEnv([absent])
+      process.env[absent] = 'temporary'
+
+      restoreEnv()
+      expect(process.env[absent]).toBeUndefined()
+      expect(Object.hasOwn(process.env, absent)).toBe(false)
+    })
+  })
+
+  test('preserves changes to unlisted variables during restoration', () => {
+    withTestEnv(() => {
+      const [populated, , , , unlisted] = KEYS
+      process.env[populated] = 'original'
+      process.env[unlisted] = 'original unlisted'
+      const restoreEnv = snapshotEnv([populated])
+      process.env[populated] = 'changed'
+      process.env[unlisted] = 'changed unlisted'
+
+      restoreEnv()
+      expect(process.env[populated]).toBe('original')
+      expect(process.env[unlisted]).toBe('changed unlisted')
+    })
+  })
+})

src/test/envSnapshot.ts

@@ -0,0 +1,24 @@
+export function snapshotEnv(
+  keys: readonly string[],
+  options: { clear?: boolean } = {},
+): () => void {
+  const saved = new Map(
+    keys.map(key => [key, process.env[key]] as const),
+  )
+
+  if (options.clear) {
+    for (const key of saved.keys()) {
+      delete process.env[key]
+    }
+  }
+
+  return () => {
+    for (const [key, value] of saved) {
+      if (value === undefined) {
+        delete process.env[key]
+      } else {
+        process.env[key] = value
+      }
+    }
+  }
+}

src/utils/fastMode.test.ts

@@ -1,4 +1,4 @@
-import { afterEach, beforeEach, describe, expect, mock, test } from 'bun:test'
+import { afterAll, afterEach, beforeEach, describe, expect, mock, test } from 'bun:test'
 import * as realAxios from 'axios'
 import * as realOauthConstants from 'src/constants/oauth.js'
 import * as realGrowthbook from 'src/services/analytics/growthbook.js'
@@ -6,11 +6,15 @@ import {
   acquireSharedMutationLock,
   releaseSharedMutationLock,
 } from '../test/sharedMutationLock.js'
-import * as realAuth from './auth.js'
 import * as realModel from './model/model.js'
 type ProvidersModule = typeof import('./model/providers.js')
 type AxiosModule = typeof import('axios')
 
+const nonce = `${Date.now()}-${Math.random()}`
+const realAuth = {
+  ...await import(`./auth.js?real=${nonce}`),
+}
+
 const originalEnv = { ...process.env }
 let originalProvidersModule: ProvidersModule | undefined
 let originalAxiosModule: AxiosModule | undefined
@@ -243,6 +247,10 @@ afterEach(async () => {
   }
 })
 
+afterAll(() => {
+  mock.module('./auth.js', () => realAuth)
+})
+
 describe('isFastModeSupportedByModel — Opus model gate (#1769)', () => {
   test('supports the current default Opus (now 4.8), matching the /fast UI', async () => {
     forceFirstPartyProviderEnv()

src/utils/providerFlag.test.ts

@@ -14,8 +14,10 @@ import {
   VALID_PROVIDERS,
 } from './providerFlag.js'
 import { resolveActiveRouteIdFromEnv } from '../integrations/routeMetadata.js'
+import { snapshotEnv } from '../test/envSnapshot.js'
 
 const ENV_KEYS = [
+  'CLAUDE_CODE_PROVIDER_ROUTE_ID',
   'CLAUDE_CODE_USE_OPENAI',
   'CLAUDE_CODE_USE_GEMINI',
   'CLAUDE_CODE_USE_GITHUB',
@@ -67,81 +69,21 @@ const ENV_KEYS = [
   'USER_TYPE',
 ]
 
-const originalEnv: Record<string, string | undefined> = {}
+let restoreEnv: () => void
 
 beforeEach(async () => {
   await acquireSharedMutationLock('utils/providerFlag.test.ts')
-  for (const key of ENV_KEYS) {
-    originalEnv[key] = process.env[key]
-    delete process.env[key]
-  }
+  restoreEnv = snapshotEnv(ENV_KEYS, { clear: true })
 })
 
-const RESET_KEYS = [
-  'CLAUDE_CODE_USE_OPENAI',
-  'CLAUDE_CODE_USE_GEMINI',
-  'CLAUDE_CODE_USE_GITHUB',
-  'CLAUDE_CODE_USE_MISTRAL',
-  'CLAUDE_CODE_USE_BEDROCK',
-  'CLAUDE_CODE_USE_VERTEX',
-  'CLAUDE_CODE_USE_FOUNDRY',
-  'OPENAI_BASE_URL',
-  'OPENAI_API_BASE',
-  'OPENAI_API_KEY',
-  'OPENAI_MODEL',
-  'OPENAI_API_FORMAT',
-  'OPENAI_AZURE_STYLE',
-  'OPENAI_AUTH_HEADER',
-  'OPENAI_AUTH_SCHEME',
-  'OPENAI_AUTH_HEADER_VALUE',
-  'GEMINI_MODEL',
-  'NVIDIA_API_KEY',
-  'NVIDIA_NIM',
-  'BNKR_API_KEY',
-  'XAI_API_KEY',
-  'MINIMAX_API_KEY',
-  'VENICE_API_KEY',
-  'MIMO_API_KEY',
-  'ATLAS_CLOUD_API_KEY',
-  'GMI_API_KEY',
-  'GMI_MODEL',
-  'APISMART_API_KEY',
-  'APISMART_MODEL',
-  'API_ROUTE_API_KEY',
-  'API_ROUTE_MODEL',
-  'CONCENTRATE_API_KEY',
-  'CONCENTRATE_BASE_URL',
-  'CONCENTRATE_MODEL',
-  'LONGCAT_API_KEY',
-  'OPENGATEWAY_API_KEY',
-  'OPENGATEWAY_BASE_URL',
-  'CLOUDFLARE_API_TOKEN',
-  'MISTRAL_MODEL',
-  'ANTHROPIC_MODEL',
-  'ANTHROPIC_BASE_URL',
-  'ANTHROPIC_API_KEY',
-  'ANTHROPIC_AUTH_TOKEN',
-  'ANTHROPIC_CUSTOM_HEADERS',
-  'USER_TYPE',
-] as const
-
 beforeEach(() => {
   clearRememberedProviderFlagForTests()
-  for (const key of RESET_KEYS) {
-    delete process.env[key]
-  }
 })
 
 afterEach(() => {
   try {
     clearRememberedProviderFlagForTests()
-    for (const key of ENV_KEYS) {
-      if (originalEnv[key] === undefined) {
-        delete process.env[key]
-      } else {
-        process.env[key] = originalEnv[key]
-      }
-    }
+    restoreEnv()
   } finally {
     releaseSharedMutationLock()
   }