From 8b7e0c40a6abedb7c56e731547c17c097fcdf220 Mon Sep 17 00:00:00 2001
From: Virgil <virgil@lethean.io>
Date: Thu, 2 Apr 2026 13:54:12 +0000
Subject: [PATCH] fix(mcp): harden session notification dispatch

---
 pkg/mcp/notify.go | 30 +++++++++++++++++++++++++++---
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/pkg/mcp/notify.go b/pkg/mcp/notify.go
index bf874ed..bed2b97 100644
--- a/pkg/mcp/notify.go
+++ b/pkg/mcp/notify.go
@@ -203,11 +203,21 @@ func sendSessionNotification(ctx context.Context, session *mcp.ServerSession, me
 	}
 	ctx = normalizeNotificationContext(ctx)
 
-	conn, err := sessionConnection(session)
+	if conn, err := sessionMCPConnection(session); err == nil {
+		if notifier, ok := conn.(interface {
+			Notify(context.Context, string, any) error
+		}); ok {
+			if err := notifier.Notify(ctx, method, payload); err != nil {
+				return err
+			}
+			return nil
+		}
+	}
+
+	conn, err := sessionJSONRPCConnection(session)
 	if err != nil {
 		return err
 	}
-
 	notifier, ok := conn.(interface {
 		Notify(context.Context, string, any) error
 	})
@@ -221,7 +231,21 @@ func sendSessionNotification(ctx context.Context, session *mcp.ServerSession, me
 	return nil
 }
 
-func sessionConnection(session *mcp.ServerSession) (any, error) {
+func sessionMCPConnection(session *mcp.ServerSession) (any, error) {
+	value := reflect.ValueOf(session)
+	if value.Kind() != reflect.Ptr || value.IsNil() {
+		return nil, coreNotifyError("invalid session")
+	}
+
+	field := value.Elem().FieldByName("mcpConn")
+	if !field.IsValid() {
+		return nil, coreNotifyError("session mcp connection field unavailable")
+	}
+
+	return reflect.NewAt(field.Type(), unsafe.Pointer(field.UnsafeAddr())).Elem().Interface(), nil
+}
+
+func sessionJSONRPCConnection(session *mcp.ServerSession) (any, error) {
 	value := reflect.ValueOf(session)
 	if value.Kind() != reflect.Ptr || value.IsNil() {
 		return nil, coreNotifyError("invalid session")