From 6f14034b6d022b963de7fec046bbd268f1083928 Mon Sep 17 00:00:00 2001
From: "Jeremy D. Miller" <jeremydmiller@yahoo.com>
Date: Sat, 28 Mar 2026 09:30:23 -0500
Subject: [PATCH] Fix FetchForWriting InvalidCastException with
 UseIdentityMapForAggregates and strongly typed IDs

Port of JasperFx/marten#4216. When UseIdentityMapForAggregates is true
and the aggregate uses a strongly typed ID (e.g., PaymentId wrapping
Guid), the identity map dictionary key type mismatch caused an
InvalidCastException. The fix makes StoreAggregateInIdentityMap
gracefully skip when the existing dictionary has an incompatible key
type.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../Bug_4214_identity_map_strong_typed_ids.cs | 100 ++++++++++++++++++
 src/Polecat/Internal/QuerySession.cs          |   9 +-
 2 files changed, 108 insertions(+), 1 deletion(-)
 create mode 100644 src/Polecat.Tests/Events/Bug_4214_identity_map_strong_typed_ids.cs

diff --git a/src/Polecat.Tests/Events/Bug_4214_identity_map_strong_typed_ids.cs b/src/Polecat.Tests/Events/Bug_4214_identity_map_strong_typed_ids.cs
new file mode 100644
index 0000000..39a52f8
--- /dev/null
+++ b/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);
+    }
+}
diff --git a/src/Polecat/Internal/QuerySession.cs b/src/Polecat/Internal/QuerySession.cs
index 9202419..8c4f69b 100644
--- a/src/Polecat/Internal/QuerySession.cs
+++ b/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
         {