fix(browser): restore replaced cookies through CDP identities (#14383)

* fix(browser): restore replaced cookies through CDP identities

Both remaining callers of the imported-domain replacement rolled back by
rebuilding cookies with cookies.set, which silently drops partitionKey.
The rollback in importValidatedCookies puts back the user's ORIGINAL
cookies that the import already deleted, so a CHIPS cookie came back as
an ordinary one and no restart recovered it.

Snapshot CDP identities before the first removal and undo through them,
the same machinery removeTransplantableCookies already uses. The store
type omits 'set' so the lossy reconstruction cannot be reintroduced, and
restoreImportedDomainCookies is deleted now that both callers are gone.

* fix(browser): skip the CDP rollback when nothing was replaced

restoreClearIdentities attaches the debugger before it iterates, so an
empty restore set would spin up a hidden BrowserWindow to put nothing
back. The old cookies.set restore was a no-op loop in that case.

src/main/browser/browser-cookie-import-policy.test.ts

@@ -3,7 +3,8 @@ import { DatabaseSync } from 'node:sqlite'
 import type { Cookie } from 'electron'
 import {
   identitiesFromClearCookies,
-  removeTransplantableCookies
+  removeTransplantableCookies,
+  type CookieClearIdentity
 } from './browser-cookie-import-clear'
 import {
   isGoogleSourceBoundCookie,
@@ -49,6 +50,29 @@ describe('isGoogleSourceBoundCookie', () => {
   })
 })
 
+// Why: mirrors what openCookieClearStore returns — get/remove plus the CDP identity pair, and
+// deliberately no 'set', so a partition-dropping reconstruction cannot be written against it.
+function replaceStore(
+  existing: Cookie[],
+  overrides: {
+    remove?: Mock
+    snapshot?: Mock
+    restore?: Mock
+  } = {}
+) {
+  const get = vi.fn().mockResolvedValue(existing)
+  const remove = overrides.remove ?? vi.fn().mockResolvedValue(undefined)
+  const snapshotClearIdentities =
+    overrides.snapshot ??
+    vi
+      .fn()
+      .mockImplementation(async (cookies: readonly { cookie: Cookie; url: string }[]) =>
+        identitiesFromClearCookies(cookies)
+      )
+  const restoreClearIdentities = overrides.restore ?? vi.fn().mockResolvedValue(undefined)
+  return { get, remove, snapshotClearIdentities, restoreClearIdentities }
+}
+
 describe('replaceCookiesForImportedDomains', () => {
   it('removes parent, exact, and child-domain cookies while preserving unrelated sites', async () => {
     const existing = [
@@ -59,99 +83,129 @@ describe('replaceCookiesForImportedDomains', () => {
       cookie('.google.com.evil.example', 'suffix-confusion'),
       cookie('.example.com', 'unrelated')
     ]
-    const get = vi.fn().mockResolvedValue(existing)
-    const remove = vi.fn().mockResolvedValue(undefined)
-    const set = vi.fn().mockResolvedValue(undefined)
+    const store = replaceStore(existing)
 
-    const removed = await replaceCookiesForImportedDomains({ get, remove, set }, [
-      'accounts.google.com'
-    ])
+    const { removed } = await replaceCookiesForImportedDomains(store, ['accounts.google.com'])
 
     expect(removed).toHaveLength(3)
-    expect(get).toHaveBeenCalledWith({})
-    expect(remove.mock.calls).toEqual([
+    expect(store.get).toHaveBeenCalledWith({})
+    expect(store.remove.mock.calls).toEqual([
       ['https://google.com/', 'parent'],
       ['https://accounts.google.com/signin', 'exact'],
       ['http://child.accounts.google.com/nested', 'child']
     ])
-    expect(set).not.toHaveBeenCalled()
+    // Why: the snapshot must cover the whole removal plan before the first removal runs.
+    expect(store.snapshotClearIdentities).toHaveBeenCalledOnce()
+    expect(
+      store.snapshotClearIdentities.mock.calls[0]?.[0].map(
+        ({ cookie: entry }: { cookie: Cookie }) => entry.name
+      )
+    ).toEqual(['parent', 'exact', 'child'])
+    expect(store.restoreClearIdentities).not.toHaveBeenCalled()
   })
 
   it('does not replace a private-suffix host cookie for a tenant import', async () => {
-    const get = vi
-      .fn()
-      .mockResolvedValue([
-        { ...cookie('github.io', 'host-only-suffix'), hostOnly: true },
-        cookie('.user.github.io', 'tenant')
-      ])
-    const remove = vi.fn().mockResolvedValue(undefined)
-    const set = vi.fn().mockResolvedValue(undefined)
+    const store = replaceStore([
+      { ...cookie('github.io', 'host-only-suffix'), hostOnly: true },
+      cookie('.user.github.io', 'tenant')
+    ])
 
-    const removed = await replaceCookiesForImportedDomains({ get, remove, set }, ['user.github.io'])
+    const { removed } = await replaceCookiesForImportedDomains(store, ['user.github.io'])
 
     expect(removed.map(({ name }) => name)).toEqual(['tenant'])
-    expect(remove).toHaveBeenCalledWith('https://user.github.io/', 'tenant')
+    expect(store.remove).toHaveBeenCalledWith('https://user.github.io/', 'tenant')
   })
 
   it('does not read or mutate the store when no valid domain scope exists', async () => {
-    const get = vi.fn()
-    const remove = vi.fn()
-    const set = vi.fn()
+    const store = replaceStore([])
 
     await expect(
-      replaceCookiesForImportedDomains({ get, remove, set }, [
-        '',
-        '...',
-        'com',
-        'co.uk',
-        'github.io'
-      ])
-    ).resolves.toEqual([])
-    expect(get).not.toHaveBeenCalled()
-    expect(remove).not.toHaveBeenCalled()
-    expect(set).not.toHaveBeenCalled()
+      replaceCookiesForImportedDomains(store, ['', '...', 'com', 'co.uk', 'github.io'])
+    ).resolves.toEqual({ removed: [], identities: [] })
+    expect(store.get).not.toHaveBeenCalled()
+    expect(store.remove).not.toHaveBeenCalled()
+    expect(store.snapshotClearIdentities).not.toHaveBeenCalled()
+    expect(store.restoreClearIdentities).not.toHaveBeenCalled()
   })
 
   it('keeps single-label intranet scopes from selecting descendant hosts', async () => {
-    const get = vi
-      .fn()
-      .mockResolvedValue([cookie('local', 'exact'), cookie('.service.local', 'descendant')])
-    const remove = vi.fn().mockResolvedValue(undefined)
-    const set = vi.fn().mockResolvedValue(undefined)
+    const store = replaceStore([cookie('local', 'exact'), cookie('.service.local', 'descendant')])
 
-    const removed = await replaceCookiesForImportedDomains({ get, remove, set }, ['local'])
+    const { removed } = await replaceCookiesForImportedDomains(store, ['local'])
 
     expect(removed.map(({ name }) => name)).toEqual(['exact'])
-    expect(remove).toHaveBeenCalledOnce()
-    expect(remove).toHaveBeenCalledWith('https://local/', 'exact')
+    expect(store.remove).toHaveBeenCalledOnce()
+    expect(store.remove).toHaveBeenCalledWith('https://local/', 'exact')
   })
 
-  it('restores cookies removed before a later removal fails', async () => {
+  // Why (STA-4097): cookies.get strips partitionKey and cookies.set ignores it, so a rollback
+  // that rebuilds cookies through the Electron API silently downgrades CHIPS cookies. The undo
+  // has to travel back through the CDP identities that actually carry the partition.
+  it('restores removed cookies through CDP identities, keeping partition identity', async () => {
     const existing = [
       cookie('.example.com', 'first', '/one'),
       cookie('.example.com', 'second', '/two')
     ]
-    const get = vi.fn().mockResolvedValue(existing)
+    const partitionKey = { topLevelSite: 'https://top.example', hasCrossSiteAncestor: true }
+    const snapshot = vi
+      .fn()
+      .mockImplementation(async (cookies: readonly { cookie: Cookie; url: string }[]) =>
+        identitiesFromClearCookies(cookies).map((identity) =>
+          identity.name === 'first' ? { ...identity, partitionKey } : identity
+        )
+      )
     const remove = vi
       .fn()
       .mockResolvedValueOnce(undefined)
       .mockRejectedValueOnce(new Error('cookie store unavailable'))
-    const set = vi.fn().mockResolvedValue(undefined)
+    const store = replaceStore(existing, { remove, snapshot })
 
-    await expect(
-      replaceCookiesForImportedDomains({ get, remove, set }, ['example.com'])
-    ).rejects.toThrow('cookie store unavailable')
-    expect(set).toHaveBeenCalledOnce()
-    expect(set).toHaveBeenCalledWith({
-      url: 'https://example.com/one',
-      name: 'first',
-      value: 'secret',
-      domain: '.example.com',
-      path: '/one',
-      secure: true,
-      httpOnly: undefined,
-      sameSite: 'unspecified'
-    })
+    await expect(replaceCookiesForImportedDomains(store, ['example.com'])).rejects.toThrow(
+      'cookie store unavailable'
+    )
+
+    expect(store.restoreClearIdentities).toHaveBeenCalledOnce()
+    const restored = store.restoreClearIdentities.mock.calls[0]?.[0] as CookieClearIdentity[]
+    // Why: the failing coordinate is restored too — a rejected remove cannot prove it survived.
+    expect(restored.map(({ name }) => name)).toEqual(['second', 'first'])
+    expect(restored.find(({ name }) => name === 'first')?.partitionKey).toEqual(partitionKey)
+    expect(restored.find(({ name }) => name === 'first')?.url).toBe('https://example.com/one')
+  })
+
+  it('aborts without removing anything when the snapshot cannot cover the removal plan', async () => {
+    const existing = [
+      cookie('.example.com', 'first', '/one'),
+      cookie('.example.com', 'second', '/two')
+    ]
+    const snapshot = vi
+      .fn()
+      .mockImplementation(async (cookies: readonly { cookie: Cookie; url: string }[]) =>
+        identitiesFromClearCookies(cookies).filter((identity) => identity.name !== 'second')
+      )
+    const store = replaceStore(existing, { snapshot })
+
+    await expect(replaceCookiesForImportedDomains(store, ['example.com'])).rejects.toThrow(
+      'the session was left unchanged'
+    )
+    expect(store.remove).not.toHaveBeenCalled()
+    expect(store.restoreClearIdentities).not.toHaveBeenCalled()
+  })
+
+  it('reports both failures when the CDP rollback itself fails', async () => {
+    const remove = vi
+      .fn()
+      .mockResolvedValueOnce(undefined)
+      .mockRejectedValueOnce(new Error('cookie store unavailable'))
+    const restore = vi.fn().mockRejectedValue(new Error('debugger detached'))
+    const store = replaceStore(
+      [cookie('.example.com', 'first', '/one'), cookie('.example.com', 'second', '/two')],
+      { remove, restore }
+    )
+
+    await expect(replaceCookiesForImportedDomains(store, ['example.com'])).rejects.toThrow(
+      'Cookie replacement and rollback failed'
+    )
+    expect(restore).toHaveBeenCalledOnce()
   })
 })
 

src/main/browser/browser-cookie-import-policy.ts

@@ -1,5 +1,7 @@
 import type { Cookie, Cookies } from 'electron'
 import { parse as parseDomain } from 'psl'
+// Why: type-only, so this does not create a runtime cycle with the clear module.
+import type { CookieClearIdentity } from './browser-cookie-import-clear'
 
 const GOOGLE_SOURCE_BOUND_COOKIE_NAMES = new Set([
   'SIDCC',
@@ -167,69 +169,102 @@ export function cookieRemovalUrl(cookie: Cookie, domain: string): string | null
   }
 }
 
-export async function restoreImportedDomainCookies(
-  store: Pick<Cookies, 'set'>,
-  cookies: readonly Cookie[]
-): Promise<void> {
-  const failures: unknown[] = []
-  for (const cookie of cookies) {
-    try {
-      const domain = cookie.domain ? normalizeCookieDomain(cookie.domain) : null
-      const url = domain ? cookieRemovalUrl(cookie, domain) : null
-      if (!url) {
-        continue
-      }
-      await store.set({
-        url,
-        name: cookie.name,
-        value: cookie.value,
-        ...(cookie.hostOnly ? {} : { domain: cookie.domain }),
-        ...(cookie.path ? { path: cookie.path } : {}),
-        secure: cookie.secure,
-        httpOnly: cookie.httpOnly,
-        sameSite: cookie.sameSite,
-        ...(cookie.expirationDate ? { expirationDate: cookie.expirationDate } : {})
-      })
-    } catch (err) {
-      failures.push(err)
+// Why (STA-4097): 'set' stays out so the partition-dropping reconstruction cannot return.
+// Undoing a removal is only possible through CDP identities, which carry partitionKey.
+export type ImportedDomainReplaceStore = Pick<Cookies, 'get' | 'remove'> & {
+  snapshotClearIdentities(
+    cookies: readonly { cookie: Cookie; url: string }[]
+  ): Promise<CookieClearIdentity[]>
+  restoreClearIdentities(identities: readonly CookieClearIdentity[]): Promise<void>
+}
+
+export type ReplacedImportedDomainCookies = {
+  removed: Cookie[]
+  identities: CookieClearIdentity[]
+}
+
+function replaceRemovalKey(url: string, name: string): string {
+  return JSON.stringify([url, name])
+}
+
+function assertIdentitiesCoverRemovable(
+  removable: readonly { cookie: Cookie; url: string }[],
+  identities: readonly CookieClearIdentity[]
+): void {
+  const covered = new Set(
+    identities.map((identity) => replaceRemovalKey(identity.url, identity.name))
+  )
+  for (const item of removable) {
+    if (!covered.has(replaceRemovalKey(item.url, item.cookie.name))) {
+      throw new Error('Could not replace existing cookies; the session was left unchanged')
     }
   }
-  if (failures.length > 0) {
-    throw new AggregateError(failures, 'Could not restore replaced cookies')
-  }
 }
 
 export async function replaceCookiesForImportedDomains(
-  store: Pick<Cookies, 'get' | 'remove' | 'set'>,
+  store: ImportedDomainReplaceStore,
   importedDomains: readonly string[]
-): Promise<Cookie[]> {
+): Promise<ReplacedImportedDomainCookies> {
   const scopes = importedDomainScopes(importedDomains)
   if (scopes.exact.size === 0) {
-    return []
+    return { removed: [], identities: [] }
   }
 
+  // Why (STA-4170): the removal plan is fixed here, beside the identities that can undo it, so
+  // the restorable set always equals the mutated set. Re-reading the jar later would widen the
+  // removal past what the snapshot can restore.
   const existingCookies = await store.get({})
-  const removedCookies: Cookie[] = []
+  const removable: { cookie: Cookie; url: string }[] = []
   for (const cookie of existingCookies) {
     const domain = cookie.domain ? normalizeCookieDomain(cookie.domain) : null
     if (!domain || !overlapsImportedDomain(cookie, domain, scopes)) {
       continue
     }
     const url = cookieRemovalUrl(cookie, domain)
-    if (!url) {
-      continue
+    if (url) {
+      removable.push({ cookie, url })
+    }
+  }
+  if (removable.length === 0) {
+    return { removed: [], identities: [] }
+  }
+
+  // Why: snapshotting before the first removal is what makes the rollback lossless; an
+  // incomplete snapshot aborts while the session is still untouched.
+  const identities = await store.snapshotClearIdentities(removable)
+  assertIdentitiesCoverRemovable(removable, identities)
+  const identitiesByKey = new Map<string, CookieClearIdentity[]>()
+  for (const identity of identities) {
+    const key = replaceRemovalKey(identity.url, identity.name)
+    const group = identitiesByKey.get(key) ?? []
+    group.push(identity)
+    identitiesByKey.set(key, group)
+  }
+
+  const removed: Cookie[] = []
+  // Why: one remove(url, name) deletes every cookie at that coordinate, partitioned twins
+  // included, so the rollback set is tracked per coordinate rather than per cookie.
+  const attemptedKeys = new Set<string>()
+  const attemptedIdentities: CookieClearIdentity[] = []
+  for (const { cookie, url } of removable) {
+    const key = replaceRemovalKey(url, cookie.name)
+    if (!attemptedKeys.has(key)) {
+      attemptedKeys.add(key)
+      attemptedIdentities.push(...(identitiesByKey.get(key) ?? []))
     }
     try {
       await store.remove(url, cookie.name)
-      removedCookies.push(cookie)
+      removed.push(cookie)
     } catch (err) {
       try {
-        await restoreImportedDomainCookies(store, removedCookies)
+        // Why: the failing coordinate is included because a rejected remove cannot prove the
+        // cookie survived; restoring a live cookie rewrites the value it was snapshotted with.
+        await store.restoreClearIdentities(attemptedIdentities.toReversed())
       } catch (restoreError) {
         throw new AggregateError([err, restoreError], 'Cookie replacement and rollback failed')
       }
       throw err
     }
   }
-  return removedCookies
+  return { removed, identities }
 }

src/main/browser/browser-cookie-import-replace-partition-rollback.electron.test.ts

@@ -0,0 +1,293 @@
+import { spawnSync } from 'node:child_process'
+import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from 'node:fs'
+import { createRequire } from 'node:module'
+import { tmpdir } from 'node:os'
+import { join } from 'node:path'
+import { afterAll, describe, expect, it } from 'vitest'
+import { build as buildVite } from 'vite'
+
+const electronBinary = createRequire(import.meta.url)('electron') as string
+const fixtureRoots: string[] = []
+
+afterAll(() => {
+  for (const root of fixtureRoots) {
+    rmSync(root, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 })
+  }
+})
+
+type CdpCookie = { name: string; value: string; partitionKey?: Record<string, unknown> }
+
+type FixtureResult = {
+  beforePartitionKey: Record<string, unknown> | undefined
+  electronGetSawPartitionKey: boolean
+  removeCalls: number
+  removeCallsBeforeThrow: number
+  snapshotCalls: number
+  restoreCalls: number
+  importedRemoveCalls: number
+  replaceError: string | null
+  chipsAfterReplace: CdpCookie[]
+  chipsAfter: CdpCookie[]
+  plainAfter: CdpCookie[]
+}
+
+const EXPECTED_PARTITION_KEY = {
+  topLevelSite: 'https://top.example',
+  hasCrossSiteAncestor: true
+}
+
+type FixtureMode = 'site-a' | 'site-b'
+
+function buildFixtureMain(bundlePath: string, resultPath: string, mode: FixtureMode): string {
+  return `
+const { app, BrowserWindow, session } = require('electron')
+const { writeFileSync } = require('node:fs')
+const { openCookieClearStore, replaceCookiesForImportedDomains } = require(${JSON.stringify(bundlePath)})
+const resultPath = ${JSON.stringify(resultPath)}
+const mode = ${JSON.stringify(mode)}
+const expectedPartitionKey = ${JSON.stringify(EXPECTED_PARTITION_KEY)}
+let currentStep = 'starting'
+const mark = (step) => {
+  currentStep = step
+  writeFileSync(resultPath, JSON.stringify({ step }))
+}
+
+const project = (cookies, name) => cookies
+  .filter((cookie) => cookie.name === name)
+  .map((cookie) => ({ name: cookie.name, value: cookie.value, partitionKey: cookie.partitionKey }))
+
+async function run() {
+  const timeout = setTimeout(() => {
+    writeFileSync(resultPath, JSON.stringify({ step: 'timed out after ' + currentStep }))
+    app.exit(1)
+  }, 20000)
+  await app.whenReady()
+  mark('ready')
+  const partition = 'persist:replace-partition-rollback-' + mode
+  const targetSession = session.fromPartition(partition)
+  const window = new BrowserWindow({ show: false, webPreferences: { partition } })
+  mark('window created')
+  await window.loadURL('data:text/html,<title>replace partition rollback fixture</title>')
+  mark('window loaded')
+  const debug = window.webContents.debugger
+  debug.attach('1.3')
+  mark('debugger attached')
+
+  // Only CDP can create a partitioned cookie; Electron's cookies API has no partitionKey.
+  await debug.sendCommand('Network.setCookie', {
+    url: 'https://app.acme-chips.test/',
+    name: 'chips-auth',
+    value: 'keep-me',
+    secure: true,
+    sameSite: 'None',
+    partitionKey: expectedPartitionKey
+  })
+  const beforeChips = (await debug.sendCommand('Network.getAllCookies')).cookies
+    .find((cookie) => cookie.name === 'chips-auth')
+  if (!beforeChips || !beforeChips.partitionKey) {
+    throw new Error('CHIPS fixture cookie was not stored partitioned')
+  }
+  mark('partitioned cookie set')
+
+  await targetSession.cookies.set({
+    url: 'https://plain.acme-chips.test/',
+    name: 'plain',
+    value: 'original',
+    secure: true
+  })
+  await targetSession.cookies.set({
+    url: 'https://victim.acme-chips.test/',
+    name: 'victim',
+    value: 'original',
+    secure: true
+  })
+  mark('ordinary cookies set')
+
+  let removeCalls = 0
+  let removeCallsBeforeThrow = 0
+  let snapshotCalls = 0
+  let restoreCalls = 0
+  let importedRemoveCalls = 0
+  let chipsRemoved = false
+  let electronGetSawPartitionKey = false
+  let replaceError = null
+  let chipsAfterReplace = []
+
+  const cookieClearStore = openCookieClearStore(targetSession)
+  const store = {
+    get: async (filter) => {
+      const cookies = await targetSession.cookies.get(filter)
+      // Why: proves the Electron API cannot see the partition it is about to destroy.
+      electronGetSawPartitionKey = cookies.some(
+        (cookie) => cookie.name === 'chips-auth' && cookie.partitionKey !== undefined
+      )
+      // Why: cookies.get order is not a contract; put CHIPS first so it is provably already
+      // removed by the time a later removal rejects.
+      return [...cookies].sort((left, right) => {
+        if (left.name === 'chips-auth') return -1
+        if (right.name === 'chips-auth') return 1
+        return 0
+      })
+    },
+    remove: async (url, name) => {
+      removeCalls++
+      if (mode === 'site-a' && chipsRemoved) {
+        removeCallsBeforeThrow = removeCalls - 1
+        throw new Error('forced removal failure after chips was removed')
+      }
+      await targetSession.cookies.remove(url, name)
+      if (name === 'chips-auth') chipsRemoved = true
+    },
+    snapshotClearIdentities: async (cookies) => {
+      snapshotCalls++
+      return cookieClearStore.snapshotClearIdentities(cookies)
+    },
+    restoreClearIdentities: async (identities) => {
+      restoreCalls++
+      return cookieClearStore.restoreClearIdentities(identities)
+    }
+  }
+
+  try {
+    if (mode === 'site-a') {
+      try {
+        await replaceCookiesForImportedDomains(store, ['acme-chips.test'])
+      } catch (error) {
+        replaceError = String(error && error.message ? error.message : error)
+      }
+      mark('site-a replace finished')
+    } else {
+      const replaced = await replaceCookiesForImportedDomains(store, ['acme-chips.test'])
+      chipsAfterReplace = project((await debug.sendCommand('Network.getAllCookies')).cookies, 'chips-auth')
+      mark('site-b replace finished')
+      // Why: stands in for importValidatedCookies' mid-import cookies.set rejection — the
+      // imported cookie is removed (lossless) and the user's originals go back through CDP.
+      await targetSession.cookies.set({
+        url: 'https://imported.acme-chips.test/',
+        name: 'imported-new',
+        value: 'new',
+        secure: true
+      })
+      await targetSession.cookies.remove('https://imported.acme-chips.test/', 'imported-new')
+      importedRemoveCalls = 1
+      mark('site-b imported cookies rolled back')
+      restoreCalls++
+      await cookieClearStore.restoreClearIdentities(replaced.identities.toReversed())
+      mark('site-b restore finished')
+    }
+  } finally {
+    cookieClearStore.dispose()
+  }
+
+  const afterCookies = (await debug.sendCommand('Network.getAllCookies')).cookies
+  clearTimeout(timeout)
+  writeFileSync(resultPath, JSON.stringify({
+    beforePartitionKey: beforeChips.partitionKey,
+    electronGetSawPartitionKey,
+    removeCalls,
+    removeCallsBeforeThrow,
+    snapshotCalls,
+    restoreCalls,
+    importedRemoveCalls,
+    replaceError,
+    chipsAfterReplace,
+    chipsAfter: project(afterCookies, 'chips-auth'),
+    plainAfter: project(afterCookies, 'plain')
+  }))
+  debug.detach()
+  window.destroy()
+  app.exit(0)
+}
+
+run().catch((error) => {
+  writeFileSync(resultPath, JSON.stringify({ step: currentStep, error: String(error && error.stack ? error.stack : error) }))
+  app.exit(1)
+})
+`
+}
+
+async function runFixture(mode: FixtureMode): Promise<FixtureResult> {
+  const root = mkdtempSync(join(tmpdir(), `orca-replace-partition-rollback-${mode}-`))
+  fixtureRoots.push(root)
+  const bundlePath = join(root, 'cookie-replace-rollback.cjs')
+  const bundleEntryPath = join(root, 'cookie-replace-rollback.ts')
+  const resultPath = join(root, 'result.json')
+  const fixturePath = join(root, 'main.cjs')
+  writeFileSync(
+    bundleEntryPath,
+    [
+      `export { openCookieClearStore } from ${JSON.stringify(join(process.cwd(), 'src/main/browser/browser-cookie-clear-store.ts'))}`,
+      `export { replaceCookiesForImportedDomains } from ${JSON.stringify(join(process.cwd(), 'src/main/browser/browser-cookie-import-policy.ts'))}`
+    ].join('\n')
+  )
+  await buildVite({
+    configFile: false,
+    logLevel: 'silent',
+    build: {
+      emptyOutDir: false,
+      lib: {
+        entry: bundleEntryPath,
+        formats: ['cjs'],
+        fileName: () => 'cookie-replace-rollback.cjs'
+      },
+      outDir: root,
+      target: 'node20',
+      rollupOptions: { external: ['electron', /^node:/] }
+    }
+  })
+  writeFileSync(fixturePath, buildFixtureMain(bundlePath, resultPath, mode))
+  const { ELECTRON_RUN_AS_NODE: _electronRunAsNode, ...env } = process.env
+  const electronArgs = [fixturePath, `--user-data-dir=${join(root, 'profile')}`]
+  const executable = process.platform === 'linux' ? 'xvfb-run' : electronBinary
+  const args =
+    process.platform === 'linux'
+      ? ['--auto-servernum', electronBinary, ...electronArgs, '--no-sandbox']
+      : electronArgs
+  const run = spawnSync(executable, args, {
+    encoding: 'utf8',
+    env,
+    timeout: 60_000
+  })
+  const fixtureResult = existsSync(resultPath) ? readFileSync(resultPath, 'utf8') : 'no result'
+  expect(run.error).toBeUndefined()
+  expect(run.status, `${fixtureResult}\n${run.stdout}\n${run.stderr}`).toBe(0)
+  return JSON.parse(fixtureResult) as FixtureResult
+}
+
+describe('partitioned cookies under a failed replace-imported-domains import', () => {
+  // Why (STA-4097): this rollback undoes deletions the import already made to the user's own
+  // cookies. Rebuilding them with cookies.set dropped partitionKey silently, so a CHIPS cookie
+  // came back as an ordinary one and no restart recovered it.
+  it('site A: restores a CHIPS cookie with its partition after a later removal rejects', async () => {
+    const result = await runFixture('site-a')
+
+    expect(result.beforePartitionKey).toEqual(EXPECTED_PARTITION_KEY)
+    // Why: without this the fixture could pass on a cookie that was never partitioned.
+    expect(result.electronGetSawPartitionKey).toBe(false)
+    // Why: assert the rollback was reached, not just that the jar looks right at the end.
+    expect(result.snapshotCalls).toBe(1)
+    expect(result.removeCallsBeforeThrow).toBeGreaterThanOrEqual(1)
+    expect(result.restoreCalls).toBe(1)
+    expect(result.replaceError).toContain('forced removal failure after chips was removed')
+    expect(result.chipsAfter).toEqual([
+      { name: 'chips-auth', value: 'keep-me', partitionKey: EXPECTED_PARTITION_KEY }
+    ])
+    expect(result.plainAfter).toEqual([{ name: 'plain', value: 'original' }])
+  }, 90_000)
+
+  it('site B: restores a CHIPS cookie with its partition after a mid-import set failure', async () => {
+    const result = await runFixture('site-b')
+
+    expect(result.beforePartitionKey).toEqual(EXPECTED_PARTITION_KEY)
+    expect(result.electronGetSawPartitionKey).toBe(false)
+    expect(result.snapshotCalls).toBe(1)
+    // Why: the replace really emptied the scope, so the restore below is doing the work.
+    expect(result.chipsAfterReplace).toEqual([])
+    expect(result.importedRemoveCalls).toBe(1)
+    expect(result.restoreCalls).toBe(1)
+    expect(result.chipsAfter).toEqual([
+      { name: 'chips-auth', value: 'keep-me', partitionKey: EXPECTED_PARTITION_KEY }
+    ])
+    expect(result.plainAfter).toEqual([{ name: 'plain', value: 'original' }])
+  }, 90_000)
+})

src/main/browser/browser-cookie-import-replacement.test.ts

@@ -3,15 +3,21 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 const {
   appGetPathMock,
   clearPendingCookieImportMock,
+  disposeClearStoreMock,
   execFileSyncMock,
+  restoreClearIdentitiesMock,
   sessionFromPartitionMock,
-  setPendingCookieImportMock
+  setPendingCookieImportMock,
+  snapshotClearIdentitiesMock
 } = vi.hoisted(() => ({
   appGetPathMock: vi.fn(),
   clearPendingCookieImportMock: vi.fn(),
+  disposeClearStoreMock: vi.fn(),
   execFileSyncMock: vi.fn(),
+  restoreClearIdentitiesMock: vi.fn(),
   sessionFromPartitionMock: vi.fn(),
-  setPendingCookieImportMock: vi.fn()
+  setPendingCookieImportMock: vi.fn(),
+  snapshotClearIdentitiesMock: vi.fn()
 }))
 
 vi.mock('./browser-session-registry', () => ({
@@ -26,6 +32,7 @@ vi.mock('electron', () => ({
   dialog: { showOpenDialog: vi.fn() },
   session: { fromPartition: sessionFromPartitionMock }
 }))
+// Why: snapshot/restore are spies, not no-ops, so a rollback that never ran cannot pass as one.
 vi.mock('./browser-cookie-clear-store', () => ({
   openCookieClearStore: (targetSession: {
     cookies: {
@@ -35,10 +42,9 @@ vi.mock('./browser-cookie-clear-store', () => ({
   }) => ({
     get: (filter: object) => targetSession.cookies.get(filter),
     remove: (url: string, name: string) => targetSession.cookies.remove(url, name),
-    snapshotClearIdentities: async (items: { cookie: Record<string, unknown>; url: string }[]) =>
-      items.map(({ cookie, url }) => ({ url, ...cookie })),
-    restoreClearIdentities: async () => undefined,
-    dispose: () => undefined
+    snapshotClearIdentities: snapshotClearIdentitiesMock,
+    restoreClearIdentities: restoreClearIdentitiesMock,
+    dispose: disposeClearStoreMock
   })
 }))
 
@@ -48,6 +54,16 @@ import { mkdtempSync, readdirSync, rmSync, writeFileSync } from 'node:fs'
 import { tmpdir } from 'node:os'
 import { join } from 'node:path'
 
+beforeEach(() => {
+  snapshotClearIdentitiesMock
+    .mockReset()
+    .mockImplementation(async (items: { cookie: Record<string, unknown>; url: string }[]) =>
+      items.map(({ cookie: entry, url }) => ({ url, ...entry }))
+    )
+  restoreClearIdentitiesMock.mockReset().mockResolvedValue(undefined)
+  disposeClearStoreMock.mockReset()
+})
+
 describe('validated cookie replacement', () => {
   let cookiesGetMock: ReturnType<typeof vi.fn>
   let cookiesRemoveMock: ReturnType<typeof vi.fn>
@@ -139,8 +155,17 @@ describe('validated cookie replacement', () => {
     expect(cookiesSetMock).not.toHaveBeenCalled()
   })
 
-  it('restores the previous snapshot when an incoming cookie is rejected', async () => {
+  // Why (STA-4097): this rollback puts back the user's ORIGINAL cookies, which the import had
+  // already deleted. cookies.get drops partitionKey and cookies.set ignores it, so rebuilding
+  // them through the Electron API resurrected CHIPS cookies unpartitioned. The undo now travels
+  // back through the CDP identities, which are the only thing that carries the partition.
+  it('restores the previous cookies through CDP identities when an incoming cookie is rejected', async () => {
+    const partitionKey = { topLevelSite: 'https://top.example', hasCrossSiteAncestor: true }
     cookiesGetMock.mockResolvedValue([cookie('.example.com', 'existing')])
+    snapshotClearIdentitiesMock.mockImplementation(
+      async (items: { cookie: Record<string, unknown>; url: string }[]) =>
+        items.map(({ cookie: entry, url }) => ({ url, ...entry, partitionKey }))
+    )
     cookiesSetMock
       .mockResolvedValue(undefined)
       .mockResolvedValueOnce(undefined)
@@ -157,16 +182,35 @@ describe('validated cookie replacement', () => {
       ['https://example.com/', 'existing'],
       ['https://example.com/', 'first']
     ])
-    expect(cookiesSetMock).toHaveBeenLastCalledWith({
-      url: 'https://example.com/',
-      name: 'existing',
-      value: 'old',
-      domain: '.example.com',
-      path: '/',
-      secure: true,
-      httpOnly: undefined,
-      sameSite: 'unspecified'
-    })
+    expect(restoreClearIdentitiesMock).toHaveBeenCalledOnce()
+    expect(restoreClearIdentitiesMock.mock.calls[0]?.[0]).toEqual([
+      { url: 'https://example.com/', ...cookie('.example.com', 'existing'), partitionKey }
+    ])
+    // Why: cookies.set is only ever the two imported cookies — a third call would mean the
+    // partition-dropping reconstruction came back.
+    expect(cookiesSetMock.mock.calls.map(([details]) => details.name)).toEqual(['first', 'second'])
+    expect(disposeClearStoreMock).toHaveBeenCalledOnce()
+  })
+
+  // Why: restoreClearIdentities attaches a debugger before it iterates, so calling it with an
+  // empty restore set would create a hidden BrowserWindow to put nothing back.
+  it('does not reach for CDP when the rollback has nothing to restore', async () => {
+    cookiesGetMock.mockResolvedValue([])
+    cookiesSetMock
+      .mockResolvedValue(undefined)
+      .mockResolvedValueOnce(undefined)
+      .mockRejectedValueOnce(new Error('cookie rejected'))
+    const filePath = writeCookies([
+      { domain: '.example.com', name: 'first', value: 'new', secure: true },
+      { domain: '.example.com', name: 'second', value: 'new', secure: true }
+    ])
+
+    const result = await importCookiesFromFile(filePath, 'persist:test')
+
+    expect(result.ok).toBe(false)
+    // Why: the imported cookie is still removed — that half of the rollback is lossless.
+    expect(cookiesRemoveMock.mock.calls).toEqual([['https://example.com/', 'first']])
+    expect(restoreClearIdentitiesMock).not.toHaveBeenCalled()
   })
 
   it('fails closed when existing cookies cannot be replaced', async () => {

src/main/browser/browser-cookie-import.ts

@@ -1,5 +1,5 @@
 /* eslint-disable max-lines -- Why: cookie import is one pipeline (detect → decrypt → stage → swap) that must stay together to keep encryption/schema/staging in sync. */
-import { app, type BrowserWindow, type Cookie, dialog, session } from 'electron'
+import { app, type BrowserWindow, dialog, session } from 'electron'
 import { execFileSync } from 'node:child_process'
 import { createDecipheriv, pbkdf2Sync, randomUUID } from 'node:crypto'
 import {
@@ -81,8 +81,8 @@ import {
   normalizeCookieDomain,
   normalizeCookieImportDomain,
   replaceCookiesForImportedDomains,
-  restoreImportedDomainCookies,
-  type CookieImportMode
+  type CookieImportMode,
+  type ReplacedImportedDomainCookies
 } from './browser-cookie-import-policy'
 import { removeTransplantableCookies, withCookieClearLock } from './browser-cookie-import-clear'
 import { openCookieClearStore } from './browser-cookie-clear-store'
@@ -594,97 +594,112 @@ async function importValidatedCookies(
   let importedCount = 0
   let skipped = totalInput - importableCookies.length
   const domainSet = new Set<string>()
-  let replacedCookies: Cookie[] | null = null
+  let replaced: ReplacedImportedDomainCookies | null = null
+  // Why (STA-4097): the rollback below has to put back cookies this import already deleted, and
+  // only CDP identities carry partitionKey — rebuilding them with cookies.set drops it silently.
+  const cookieClearStore =
+    mode === 'replace-imported-domains' && importableCookies.length > 0
+      ? openCookieClearStore(targetSession)
+      : null
 
-  if (mode === 'replace-imported-domains' && importableCookies.length > 0) {
-    try {
-      replacedCookies = await replaceCookiesForImportedDomains(
-        targetSession.cookies,
-        importableCookies.map((cookie) => cookie.domain)
-      )
-      diag(`  removed ${replacedCookies.length} existing cookies in imported domain scopes`)
-    } catch (err) {
-      diag(`  existing cookie replacement failed: ${summarizeCookieImportError(err)}`)
+  try {
+    if (cookieClearStore) {
+      try {
+        replaced = await replaceCookiesForImportedDomains(
+          cookieClearStore,
+          importableCookies.map((cookie) => cookie.domain)
+        )
+        diag(`  removed ${replaced.removed.length} existing cookies in imported domain scopes`)
+      } catch (err) {
+        diag(`  existing cookie replacement failed: ${summarizeCookieImportError(err)}`)
+        return {
+          ok: false,
+          reason: reasonWithDiagLog('Could not replace existing cookies for the imported sites.')
+        }
+      }
+    }
+
+    // Why: Electron's cookies.set() rejects any non-printable-ASCII byte; strip as a safety net.
+    const stripNonPrintable = (s: string): string => s.replace(/[^\x20-\x7E]/g, '')
+    const importedCookieKeys: { url: string; name: string }[] = []
+    let setFailure: unknown = null
+
+    for (const cookie of importableCookies) {
+      try {
+        // Why: Chromium rejects __Host- cookies unless they omit domain and use path=/.
+        const isHostPrefixed = cookie.name.startsWith('__Host-')
+        const path = isHostPrefixed ? '/' : cookie.path
+        await targetSession.cookies.set({
+          url: cookie.url,
+          name: cookie.name,
+          value: stripNonPrintable(cookie.value),
+          ...(isHostPrefixed ? {} : { domain: cookie.domain }),
+          path,
+          secure: cookie.secure,
+          httpOnly: cookie.httpOnly,
+          sameSite: cookie.sameSite,
+          expirationDate: cookie.expirationDate
+        })
+        const removalUrl = new URL(cookie.url)
+        removalUrl.pathname = path.startsWith('/') ? path : '/'
+        importedCookieKeys.push({ url: removalUrl.toString(), name: cookie.name })
+        importedCount++
+        // Why: surface only the domain (never name/value/path) so the summary doesn't leak secret cookie data.
+        const cleanDomain = cookie.domain.startsWith('.') ? cookie.domain.slice(1) : cookie.domain
+        domainSet.add(cleanDomain)
+      } catch (err) {
+        skipped++
+        setFailure = err
+        if (skipped <= 5) {
+          // Find the exact offending character position and code
+          const val = cookie.value
+          let badInfo = 'none found'
+          for (let i = 0; i < val.length; i++) {
+            const code = val.charCodeAt(i)
+            if (code < 0x20 || code > 0x7e) {
+              badInfo = `pos=${i} char=U+${code.toString(16).padStart(4, '0')}`
+              break
+            }
+          }
+          diag(
+            `  cookie.set FAILED: domain=${cookie.domain} name=${cookie.name} valLen=${val.length} badChar=${badInfo} err=${String(err)}`
+          )
+        }
+        if (replaced) {
+          break
+        }
+      }
+    }
+
+    // Why: replaced is only ever set alongside cookieClearStore, which owns the CDP restore.
+    if (setFailure && replaced && cookieClearStore) {
+      const rollbackFailures: unknown[] = []
+      for (const cookie of importedCookieKeys.toReversed()) {
+        try {
+          await targetSession.cookies.remove(cookie.url, cookie.name)
+        } catch (err) {
+          rollbackFailures.push(err)
+        }
+      }
+      // Why: restoreClearIdentities attaches the debugger before it iterates, so an empty
+      // restore set would spin up a hidden BrowserWindow to put nothing back.
+      if (replaced.identities.length > 0) {
+        try {
+          await cookieClearStore.restoreClearIdentities(replaced.identities.toReversed())
+        } catch (err) {
+          rollbackFailures.push(err)
+        }
+      }
+      if (rollbackFailures.length > 0) {
+        diag(`  cookie replacement rollback failed: ${rollbackFailures.length} operation(s)`)
+      }
       return {
         ok: false,
-        reason: reasonWithDiagLog('Could not replace existing cookies for the imported sites.')
+        reason: reasonWithDiagLog('Could not safely replace cookies for the imported sites.')
       }
     }
-  }
-
-  // Why: Electron's cookies.set() rejects any non-printable-ASCII byte; strip as a safety net.
-  const stripNonPrintable = (s: string): string => s.replace(/[^\x20-\x7E]/g, '')
-  const importedCookieKeys: { url: string; name: string }[] = []
-  let setFailure: unknown = null
-
-  for (const cookie of importableCookies) {
-    try {
-      // Why: Chromium rejects __Host- cookies unless they omit domain and use path=/.
-      const isHostPrefixed = cookie.name.startsWith('__Host-')
-      const path = isHostPrefixed ? '/' : cookie.path
-      await targetSession.cookies.set({
-        url: cookie.url,
-        name: cookie.name,
-        value: stripNonPrintable(cookie.value),
-        ...(isHostPrefixed ? {} : { domain: cookie.domain }),
-        path,
-        secure: cookie.secure,
-        httpOnly: cookie.httpOnly,
-        sameSite: cookie.sameSite,
-        expirationDate: cookie.expirationDate
-      })
-      const removalUrl = new URL(cookie.url)
-      removalUrl.pathname = path.startsWith('/') ? path : '/'
-      importedCookieKeys.push({ url: removalUrl.toString(), name: cookie.name })
-      importedCount++
-      // Why: surface only the domain (never name/value/path) so the summary doesn't leak secret cookie data.
-      const cleanDomain = cookie.domain.startsWith('.') ? cookie.domain.slice(1) : cookie.domain
-      domainSet.add(cleanDomain)
-    } catch (err) {
-      skipped++
-      setFailure = err
-      if (skipped <= 5) {
-        // Find the exact offending character position and code
-        const val = cookie.value
-        let badInfo = 'none found'
-        for (let i = 0; i < val.length; i++) {
-          const code = val.charCodeAt(i)
-          if (code < 0x20 || code > 0x7e) {
-            badInfo = `pos=${i} char=U+${code.toString(16).padStart(4, '0')}`
-            break
-          }
-        }
-        diag(
-          `  cookie.set FAILED: domain=${cookie.domain} name=${cookie.name} valLen=${val.length} badChar=${badInfo} err=${String(err)}`
-        )
-      }
-      if (replacedCookies) {
-        break
-      }
-    }
-  }
-
-  if (setFailure && replacedCookies) {
-    const rollbackFailures: unknown[] = []
-    for (const cookie of importedCookieKeys.toReversed()) {
-      try {
-        await targetSession.cookies.remove(cookie.url, cookie.name)
-      } catch (err) {
-        rollbackFailures.push(err)
-      }
-    }
-    try {
-      await restoreImportedDomainCookies(targetSession.cookies, replacedCookies)
-    } catch (err) {
-      rollbackFailures.push(err)
-    }
-    if (rollbackFailures.length > 0) {
-      diag(`  cookie replacement rollback failed: ${rollbackFailures.length} operation(s)`)
-    }
-    return {
-      ok: false,
-      reason: reasonWithDiagLog('Could not safely replace cookies for the imported sites.')
-    }
+  } finally {
+    cookieClearStore?.dispose()
   }
 
   diag(