fix: correctly process NoSuchUserContextException in cookie operations

src/bidiMapper/domains/storage/StorageProcessor.ts

@@ -15,7 +15,10 @@
  * limitations under the License.
  */
 import type {CdpClient} from '../../../cdp/CdpClient.js';
-import {UnableToSetCookieException} from '../../../protocol/protocol.js';
+import {
+  NoSuchUserContextException,
+  UnableToSetCookieException,
+} from '../../../protocol/protocol.js';
 import type {Storage, Network} from '../../../protocol/protocol.js';
 import {assert} from '../../../utils/assert.js';
 import type {LoggerFn} from '../../../utils/log.js';
@@ -51,12 +54,21 @@ export class StorageProcessor {
   ): Promise<Storage.DeleteCookiesResult> {
     const partitionKey = this.#expandStoragePartitionSpec(params.partition);
 
-    const cdpResponse = await this.#browserCdpClient.sendCommand(
-      'Storage.getCookies',
-      {
-        browserContextId: this.#getCdpBrowserContextId(partitionKey),
+    let cdpResponse;
+    try {
+      cdpResponse = await this.#browserCdpClient.sendCommand(
+        'Storage.getCookies',
+        {
+          browserContextId: this.#getCdpBrowserContextId(partitionKey),
+        }
+      );
+    } catch (err: any) {
+      if (this.#isNoSuchUserContextError(err)) {
+        // If the user context is not found, special error is thrown.
+        throw new NoSuchUserContextException(err.message);
       }
-    );
+      throw err;
+    }
 
     const cdpCookiesToDelete = cdpResponse.cookies
       .filter(
@@ -91,12 +103,21 @@ export class StorageProcessor {
   ): Promise<Storage.GetCookiesResult> {
     const partitionKey = this.#expandStoragePartitionSpec(params.partition);
 
-    const cdpResponse = await this.#browserCdpClient.sendCommand(
-      'Storage.getCookies',
-      {
-        browserContextId: this.#getCdpBrowserContextId(partitionKey),
+    let cdpResponse;
+    try {
+      cdpResponse = await this.#browserCdpClient.sendCommand(
+        'Storage.getCookies',
+        {
+          browserContextId: this.#getCdpBrowserContextId(partitionKey),
+        }
+      );
+    } catch (err: any) {
+      if (this.#isNoSuchUserContextError(err)) {
+        // If the user context is not found, special error is thrown.
+        throw new NoSuchUserContextException(err.message);
       }
-    );
+      throw err;
+    }
 
     const filteredBiDiCookies = cdpResponse.cookies
       .filter(
@@ -127,15 +148,26 @@ export class StorageProcessor {
         cookies: [cdpCookie],
         browserContextId: this.#getCdpBrowserContextId(partitionKey),
       });
-    } catch (e: any) {
-      this.#logger?.(LogType.debugError, e);
-      throw new UnableToSetCookieException(e.toString());
+    } catch (err: any) {
+      if (this.#isNoSuchUserContextError(err)) {
+        // If the user context is not found, special error is thrown.
+        throw new NoSuchUserContextException(err.message);
+      }
+
+      this.#logger?.(LogType.debugError, err);
+      throw new UnableToSetCookieException(err.toString());
     }
     return {
       partitionKey,
     };
   }
 
+  #isNoSuchUserContextError(err: Error): boolean {
+    // Heuristic to detect if the user context is not found.
+    // See https://source.chromium.org/chromium/chromium/src/+/main:content/browser/devtools/protocol/browser_handler.cc;drc=a56154dd81e4679712422ac6eed2c9581cb51ab0;l=314
+    return err.message?.startsWith('Failed to find browser context for id');
+  }
+
   #getCdpBrowserContextId(
     partitionKey: Storage.PartitionKey
   ): string | undefined {

tests/storage/test_delete_cookies.py

@@ -132,6 +132,28 @@ async def test_cookies_delete_partition_user_context(websocket, context_id,
     await assert_only_cookies_present(websocket, [default_user_context_cookie])
 
 
+@pytest.mark.asyncio
+async def test_cookies_delete_partition_user_context_unknown(
+        websocket, context_id):
+    user_context_partition = {
+        'type': 'storageKey',
+        'userContext': 'UNKNOWN_USER_CONTEXT',
+    }
+
+    with pytest.raises(Exception,
+                       match=str({
+                           'error': 'no such user context',
+                           'message': '.*'
+                       })):
+        await execute_command(
+            websocket, {
+                'method': 'storage.deleteCookies',
+                'params': {
+                    'partition': user_context_partition,
+                }
+            })
+
+
 @pytest.mark.asyncio
 async def test_cookies_delete_partition_unsupported_key(websocket, context_id):
     cookie = get_bidi_cookie(SOME_COOKIE_NAME, SOME_COOKIE_VALUE, SOME_DOMAIN)

tests/storage/test_get_cookies.py

@@ -149,6 +149,28 @@ async def test_cookies_get_partition_user_context(websocket, context_id,
     }
 
 
+@pytest.mark.asyncio
+async def test_cookies_get_partition_user_context_unknown(
+        websocket, context_id):
+    user_context_partition = {
+        'type': 'storageKey',
+        'userContext': 'UNKNOWN_USER_CONTEXT',
+    }
+
+    with pytest.raises(Exception,
+                       match=str({
+                           'error': 'no such user context',
+                           'message': '.*'
+                       })):
+        await execute_command(
+            websocket, {
+                'method': 'storage.getCookies',
+                'params': {
+                    'partition': user_context_partition,
+                }
+            })
+
+
 @pytest.mark.asyncio
 async def test_cookies_get_partition_unsupported_key(websocket, context_id):
     cookie = get_bidi_cookie(SOME_COOKIE_NAME, SOME_COOKIE_VALUE, SOME_DOMAIN)

tests/storage/test_set_cookies.py

@@ -147,7 +147,12 @@ async def test_cookie_set_partition_browsing_context(websocket, context_id):
 
 
 @pytest.mark.asyncio
-async def test_cookie_set_partition_user_context(websocket, context_id):
+async def test_cookie_set_partition_user_context(websocket, context_id,
+                                                 user_context_id):
+    user_context_partition = {
+        'type': 'storageKey',
+        'userContext': user_context_id,
+    }
     resp = await execute_command(
         websocket, {
             'method': 'storage.setCookie',
@@ -161,18 +166,18 @@ async def test_cookie_set_partition_user_context(websocket, context_id):
                     },
                     'domain': SOME_DOMAIN,
                 },
-                'partition': {
-                    'type': 'context',
-                    'context': context_id
-                }
+                'partition': user_context_partition
             }
         })
-    assert resp == {'partitionKey': {'userContext': 'default'}}
+    assert resp == {'partitionKey': {'userContext': user_context_id}}
 
-    resp = await execute_command(websocket, {
-        'method': 'storage.getCookies',
-        'params': {}
-    })
+    resp = await execute_command(
+        websocket, {
+            'method': 'storage.getCookies',
+            'params': {
+                'partition': user_context_partition
+            }
+        })
     assert resp == {
         'cookies': [
             AnyExtending(
@@ -182,11 +187,42 @@ async def test_cookie_set_partition_user_context(websocket, context_id):
                                 secure=True))
         ],
         'partitionKey': {
-            'userContext': 'default'
+            'userContext': user_context_id
         }
     }
 
 
+@pytest.mark.asyncio
+async def test_cookie_set_partition_user_context_unknown(
+        websocket, context_id):
+    user_context_partition = {
+        'type': 'storageKey',
+        'userContext': 'UNKNOWN_USER_CONTEXT',
+    }
+
+    with pytest.raises(Exception,
+                       match=str({
+                           'error': 'no such user context',
+                           'message': '.*'
+                       })):
+        await execute_command(
+            websocket, {
+                'method': 'storage.setCookie',
+                'params': {
+                    'cookie': {
+                        'secure': True,
+                        'name': SOME_COOKIE_NAME,
+                        'value': {
+                            'type': 'string',
+                            'value': SOME_COOKIE_VALUE
+                        },
+                        'domain': SOME_DOMAIN,
+                    },
+                    'partition': user_context_partition
+                }
+            })
+
+
 @pytest.mark.asyncio
 async def test_cookie_set_partition_browsing_context_from_user_context(
         websocket, create_context, user_context_id):