Fix FetchForWriting with UseIdentityMapForAggregates and strongly typed IDs

src/Polecat.Tests/Events/Bug_4214_identity_map_strong_typed_ids.cs

@@ -0,0 +1,100 @@
+using JasperFx.Events.Projections;
+using Polecat.Projections;
+using Polecat.Tests.Harness;
+using Polecat.Tests.Projections;
+
+namespace Polecat.Tests.Events;
+
+/// <summary>
+/// Regression test for https://github.com/JasperFx/marten/issues/4214
+/// FetchForWriting throws InvalidCastException when using UseIdentityMapForAggregates
+/// with strongly typed IDs.
+/// </summary>
+[Collection("integration")]
+public class Bug_4214_identity_map_strong_typed_ids : IntegrationContext
+{
+    public Bug_4214_identity_map_strong_typed_ids(DefaultStoreFixture fixture) : base(fixture)
+    {
+    }
+
+    [Fact]
+    public async Task fetch_for_writing_with_identity_map_and_strong_typed_guid_id_inline()
+    {
+        await StoreOptions(opts =>
+        {
+            opts.DatabaseSchemaName = "bug4214_inline";
+            opts.Projections.Add(new SingleStreamProjection<Payment, PaymentId>(), ProjectionLifecycle.Inline);
+            opts.Projections.UseIdentityMapForAggregates = true;
+        });
+
+        await using var session = theStore.LightweightSession();
+
+        var id = Guid.NewGuid();
+        session.Events.StartStream(id,
+            new PaymentCreated(DateTimeOffset.UtcNow),
+            new PaymentVerified(DateTimeOffset.UtcNow));
+
+        await session.SaveChangesAsync();
+
+        // This threw InvalidCastException before the fix
+        var stream = await session.Events.FetchForWriting<Payment>(id);
+        stream.Aggregate.ShouldNotBeNull();
+        stream.Aggregate!.State.ShouldBe(PaymentState.Verified);
+    }
+
+    [Fact]
+    public async Task fetch_for_writing_with_identity_map_and_strong_typed_guid_id_live()
+    {
+        await StoreOptions(opts =>
+        {
+            opts.DatabaseSchemaName = "bug4214_live";
+            opts.Projections.Add(new SingleStreamProjection<Payment, PaymentId>(), ProjectionLifecycle.Live);
+            opts.Projections.UseIdentityMapForAggregates = true;
+        });
+
+        await using var session = theStore.LightweightSession();
+
+        var id = Guid.NewGuid();
+        session.Events.StartStream(id,
+            new PaymentCreated(DateTimeOffset.UtcNow),
+            new PaymentVerified(DateTimeOffset.UtcNow));
+
+        await session.SaveChangesAsync();
+
+        var stream = await session.Events.FetchForWriting<Payment>(id);
+        stream.Aggregate.ShouldNotBeNull();
+        stream.Aggregate!.State.ShouldBe(PaymentState.Verified);
+    }
+
+    [Fact]
+    public async Task fetch_for_writing_twice_with_identity_map_and_strong_typed_guid_id_inline()
+    {
+        await StoreOptions(opts =>
+        {
+            opts.DatabaseSchemaName = "bug4214_twice";
+            opts.Projections.Add(new SingleStreamProjection<Payment, PaymentId>(), ProjectionLifecycle.Inline);
+            opts.Projections.UseIdentityMapForAggregates = true;
+        });
+
+        await using var session = theStore.LightweightSession();
+
+        var id = Guid.NewGuid();
+        session.Events.StartStream(id,
+            new PaymentCreated(DateTimeOffset.UtcNow),
+            new PaymentVerified(DateTimeOffset.UtcNow));
+
+        await session.SaveChangesAsync();
+
+        // First fetch stores in identity map
+        var stream1 = await session.Events.FetchForWriting<Payment>(id);
+        stream1.Aggregate.ShouldNotBeNull();
+
+        stream1.AppendOne(new PaymentCanceled(DateTimeOffset.UtcNow));
+        await session.SaveChangesAsync();
+
+        // Second fetch should retrieve from identity map without cast error
+        var stream2 = await session.Events.FetchForWriting<Payment>(id);
+        stream2.Aggregate.ShouldNotBeNull();
+        stream2.Aggregate!.State.ShouldBe(PaymentState.Canceled);
+    }
+}

src/Polecat/Internal/QuerySession.cs

@@ -36,7 +36,14 @@ internal void StoreAggregateInIdentityMap<TDoc, TId>(TId id, TDoc document)
     {
         if (AggregateIdentityMap.TryGetValue(typeof(TDoc), out var raw))
         {
-            ((Dictionary<TId, TDoc>)raw)[id] = document;
+            if (raw is Dictionary<TId, TDoc> typedDict)
+            {
+                typedDict[id] = document;
+            }
+            // else: The identity map was created with a different key type (e.g., a strong-typed ID
+            // like PaymentId while TId is Guid). The document is already stored by the inline
+            // projection under the strong-typed key, so we skip storing it again to avoid
+            // replacing the dictionary with an incompatible key type.
         }
         else
         {