Fix CookiePartitionKey serialization and add partition key test

lib/PuppeteerSharp.Nunit/TestExpectations/TestExpectations.local.json

@@ -110,5 +110,12 @@
     "platforms": ["darwin", "linux", "win32"],
     "parameters": ["webDriverBiDi"],
     "expectations": ["SKIP"]
+  },
+  {
+    "comment": "Firefox native BiDi does not return goog:partitionKey or goog:sourceScheme extension data",
+    "testIdPattern": "[cookies.spec] Cookie specs Page.setCookie should set a cookie with a partitionKey",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["firefox"],
+    "expectations": ["FAIL"]
   }
 ]

lib/PuppeteerSharp.Tests/CookiesTests/SetCookiesTests.cs

@@ -1,3 +1,4 @@
+using System;
 using System.Linq;
 using System.Threading.Tasks;
 using NUnit.Framework;
@@ -182,6 +183,50 @@ await Page.SetCookieAsync(new CookieParam
             Assert.That(cookie.Session, Is.True);
         }
 
+        [Test, PuppeteerTest("cookies.spec", "Cookie specs Page.setCookie", "should set a cookie with a partitionKey")]
+        public async Task ShouldSetACookieWithAPartitionKey()
+        {
+            var options = TestConstants.DefaultBrowserOptions();
+            options.AcceptInsecureCerts = true;
+
+            await using var browser = await Puppeteer.LaunchAsync(options, TestConstants.LoggerFactory);
+            await using var page = await browser.NewPageAsync();
+            await page.GoToAsync(TestConstants.HttpsPrefix + "/empty.html");
+            var url = new Uri(page.Url);
+            var key = url.GetLeftPart(UriPartial.Authority);
+            await page.SetCookieAsync(new CookieParam
+            {
+                Url = url.AbsoluteUri,
+                Name = "partitionCookie",
+                Value = "partition",
+                Secure = true,
+                PartitionKey = key,
+            });
+            var cookies = await page.GetCookiesAsync();
+            Assert.That(cookies, Has.Exactly(1).Items);
+            var cookie = cookies.First();
+            Assert.That(cookie.Name, Is.EqualTo("partitionCookie"));
+            Assert.That(cookie.Value, Is.EqualTo("partition"));
+            Assert.That(cookie.Domain, Is.EqualTo(url.Host));
+            Assert.That(cookie.Path, Is.EqualTo("/"));
+            Assert.That(cookie.Expires, Is.EqualTo(-1));
+            Assert.That(cookie.Size, Is.EqualTo(24));
+            Assert.That(cookie.HttpOnly, Is.False);
+            Assert.That(cookie.Secure, Is.True);
+            Assert.That(cookie.Session, Is.True);
+            if (TestConstants.IsChrome)
+            {
+                Assert.That(cookie.SourceScheme, Is.EqualTo(CookieSourceScheme.Secure));
+            }
+
+            if (TestConstants.IsChrome)
+            {
+                key = url.GetComponents(UriComponents.Scheme | UriComponents.Host, UriFormat.UriEscaped);
+            }
+
+            Assert.That(cookie.PartitionKey, Is.EqualTo(key));
+        }
+
         [Test, PuppeteerTest("cookies.spec", "Cookie specs Page.setCookie", "should not set a cookie on a blank page")]
         public async Task ShouldNotSetACookieOnABlankPage()
         {

lib/PuppeteerSharp.Tests/UtilitiesTests/CookiePartitionKeyConverterTests.cs

@@ -10,7 +10,7 @@ namespace PuppeteerSharp.Tests.UtilitiesTests
     public class CookiePartitionKeyConverterTests
     {
         [Test]
-        public void ShouldSerializeAndDeserializeAsString()
+        public void ShouldSerializeAndDeserializeAsObject()
         {
             // Arrange
             var cookie = new CookieParam
@@ -23,8 +23,8 @@ public void ShouldSerializeAndDeserializeAsString()
             // Act - Serialize to JSON
             var json = JsonSerializer.Serialize(cookie);
 
-            // Assert - Verify it's serialized as a string
-            Assert.That(json, Does.Contain("\"PartitionKey\":\"https://example.com\""));
+            // Assert - Verify it's serialized as an object with topLevelSite and hasCrossSiteAncestor
+            Assert.That(json, Does.Contain("\"PartitionKey\":{\"topLevelSite\":\"https://example.com\",\"hasCrossSiteAncestor\":false}"));
 
             // Act - Deserialize back
             var deserialized = JsonSerializer.Deserialize<CookieParam>(json);

lib/PuppeteerSharp/Bidi/BidiCookieHelper.cs

@@ -23,6 +23,7 @@
 #if !CDP_ONLY
 
 using System;
+using System.Text.Json;
 using WebDriverBiDi.Network;
 using WebDriverBiDi.Storage;
 using BidiCookie = WebDriverBiDi.Network.Cookie;
@@ -47,7 +48,8 @@ public static CookieParam BidiToPuppeteerCookie(BidiCookie bidiCookie)
             SameSite = ConvertSameSiteBidiToPuppeteer(bidiCookie.SameSite),
             Expires = bidiCookie.EpochExpires.HasValue ? (double)bidiCookie.EpochExpires.Value / 1000 : -1,
             Session = !bidiCookie.EpochExpires.HasValue || bidiCookie.EpochExpires.Value == 0,
-            SourceScheme = CookieSourceScheme.Unset,
+            SourceScheme = GetSourceScheme(bidiCookie),
+            PartitionKey = GetPartitionKey(bidiCookie),
         };
 
     /// <summary>
@@ -212,6 +214,51 @@ private static string ConvertPriorityEnumToString(CookiePriority priority)
             _ => "Medium",
         };
     }
+
+    private static string GetPartitionKey(BidiCookie bidiCookie)
+    {
+        if (!bidiCookie.AdditionalData.TryGetValue("goog:partitionKey", out var value))
+        {
+            return null;
+        }
+
+        if (value is string stringValue)
+        {
+            return stringValue;
+        }
+
+        if (value is JsonElement jsonElement)
+        {
+            if (jsonElement.ValueKind == JsonValueKind.String)
+            {
+                return jsonElement.GetString();
+            }
+
+            if (jsonElement.ValueKind == JsonValueKind.Object &&
+                jsonElement.TryGetProperty("topLevelSite", out var topLevelSite))
+            {
+                return topLevelSite.GetString();
+            }
+        }
+
+        return null;
+    }
+
+    private static CookieSourceScheme GetSourceScheme(BidiCookie bidiCookie)
+    {
+        if (!bidiCookie.AdditionalData.TryGetValue("goog:sourceScheme", out var value))
+        {
+            return CookieSourceScheme.Unset;
+        }
+
+        var scheme = value?.ToString();
+        return scheme switch
+        {
+            "Secure" => CookieSourceScheme.Secure,
+            "NonSecure" => CookieSourceScheme.NonSecure,
+            _ => CookieSourceScheme.Unset,
+        };
+    }
 }
 
 #endif

lib/PuppeteerSharp/Bidi/BidiPage.cs

@@ -716,12 +716,6 @@ public override async Task SetCookieAsync(params CookieParam[] cookies)
                 throw new PuppeteerException($"Data URL page can not have cookie \"{cookie.Name}\"");
             }
 
-            // TODO: Support Chrome cookie partition keys
-            if (!string.IsNullOrEmpty(cookie.PartitionKey))
-            {
-                throw new PuppeteerException("BiDi only allows domain partition keys");
-            }
-
             Uri normalizedUrl = null;
             if (Uri.TryCreate(cookieUrl, UriKind.Absolute, out var parsedUrl))
             {
@@ -756,10 +750,25 @@ public override async Task SetCookieAsync(params CookieParam[] cookies)
 
             var bidiCookie = BidiCookieHelper.PuppeteerToBidiCookie(cookieToUse, domain);
 
-            await BidiBrowser.Driver.Storage.SetCookieAsync(new WebDriverBiDi.Storage.SetCookieCommandParameters(bidiCookie)
+            if (!string.IsNullOrEmpty(cookie.PartitionKey))
             {
-                Partition = new WebDriverBiDi.Storage.BrowsingContextPartitionDescriptor(BidiMainFrame.BrowsingContext.Id),
-            }).ConfigureAwait(false);
+                var userContext = ((BidiBrowserContext)BrowserContext).UserContext;
+                await BidiBrowser.Driver.Storage.SetCookieAsync(new WebDriverBiDi.Storage.SetCookieCommandParameters(bidiCookie)
+                {
+                    Partition = new WebDriverBiDi.Storage.StorageKeyPartitionDescriptor
+                    {
+                        SourceOrigin = cookie.PartitionKey,
+                        UserContextId = userContext.Id,
+                    },
+                }).ConfigureAwait(false);
+            }
+            else
+            {
+                await BidiBrowser.Driver.Storage.SetCookieAsync(new WebDriverBiDi.Storage.SetCookieCommandParameters(bidiCookie)
+                {
+                    Partition = new WebDriverBiDi.Storage.BrowsingContextPartitionDescriptor(BidiMainFrame.BrowsingContext.Id),
+                }).ConfigureAwait(false);
+            }
         }
     }
 

lib/PuppeteerSharp/Helpers/Json/CookiePartitionKeyConverter.cs

@@ -35,11 +35,12 @@ public override void Write(
             string value,
             JsonSerializerOptions options)
         {
-            // Write as a simple string for user serialization/deserialization
-            // This allows cookies to be easily saved to and loaded from files
             if (value != null)
             {
-                writer.WriteStringValue(value);
+                writer.WriteStartObject();
+                writer.WriteString("topLevelSite", value);
+                writer.WriteBoolean("hasCrossSiteAncestor", false);
+                writer.WriteEndObject();
             }
             else
             {