From c3a449e678cd84281524719a41a33657c6402175 Mon Sep 17 00:00:00 2001
From: Virgil <virgil@lethean.io>
Date: Thu, 2 Apr 2026 10:12:35 +0000
Subject: [PATCH] fix(mcp): add session notification helper

Use the Core context for forwarded IPC channel events and expose a session-level notification helper that the broadcast path reuses. This keeps the notification API more symmetric and avoids dropping context during IPC forwarding.

Co-Authored-By: Virgil <virgil@lethean.io>
---
 pkg/mcp/notify.go      | 26 +++++++++++++++++++-------
 pkg/mcp/notify_test.go | 12 ++++++++++++
 pkg/mcp/register.go    |  9 ++++++++-
 3 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/pkg/mcp/notify.go b/pkg/mcp/notify.go
index 740cadf..aba9afa 100644
--- a/pkg/mcp/notify.go
+++ b/pkg/mcp/notify.go
@@ -54,13 +54,25 @@ type ChannelNotification struct {
 //	s.SendNotificationToAllClients(ctx, "info", "monitor", map[string]any{"event": "build complete"})
 func (s *Service) SendNotificationToAllClients(ctx context.Context, level mcp.LoggingLevel, logger string, data any) {
 	for session := range s.server.Sessions() {
-		if err := session.Log(ctx, &mcp.LoggingMessageParams{
-			Level:  level,
-			Logger: logger,
-			Data:   data,
-		}); err != nil {
-			s.logger.Debug("notify: failed to send to session", "session", session.ID(), "error", err)
-		}
+		s.SendNotificationToSession(ctx, session, level, logger, data)
+	}
+}
+
+// SendNotificationToSession sends a log-level notification to one connected
+// MCP session.
+//
+//	s.SendNotificationToSession(ctx, session, "info", "monitor", data)
+func (s *Service) SendNotificationToSession(ctx context.Context, session *mcp.ServerSession, level mcp.LoggingLevel, logger string, data any) {
+	if session == nil {
+		return
+	}
+
+	if err := session.Log(ctx, &mcp.LoggingMessageParams{
+		Level:  level,
+		Logger: logger,
+		Data:   data,
+	}); err != nil {
+		s.logger.Debug("notify: failed to send to session", "session", session.ID(), "error", err)
 	}
 }
 
diff --git a/pkg/mcp/notify_test.go b/pkg/mcp/notify_test.go
index a1ea4f2..def8148 100644
--- a/pkg/mcp/notify_test.go
+++ b/pkg/mcp/notify_test.go
@@ -45,6 +45,18 @@ func TestChannelSendToSession_Good_GuardNilSession(t *testing.T) {
 	})
 }
 
+func TestSendNotificationToSession_Good_GuardNilSession(t *testing.T) {
+	svc, err := New(Options{})
+	if err != nil {
+		t.Fatalf("New() failed: %v", err)
+	}
+
+	ctx := context.Background()
+	svc.SendNotificationToSession(ctx, nil, "info", "test", map[string]any{
+		"ok": true,
+	})
+}
+
 func TestChannelSendToSession_Good_CustomNotification(t *testing.T) {
 	svc, err := New(Options{})
 	if err != nil {
diff --git a/pkg/mcp/register.go b/pkg/mcp/register.go
index 8b8de8b..b7e7b92 100644
--- a/pkg/mcp/register.go
+++ b/pkg/mcp/register.go
@@ -97,9 +97,16 @@ func (s *Service) OnStartup(ctx context.Context) core.Result {
 //
 //	c.ACTION(mcp.ChannelPush{Channel: "agent.status", Data: statusMap})
 func (s *Service) HandleIPCEvents(c *core.Core, msg core.Message) core.Result {
+	ctx := context.Background()
+	if c != nil {
+		if coreCtx := c.Context(); coreCtx != nil {
+			ctx = coreCtx
+		}
+	}
+
 	switch ev := msg.(type) {
 	case ChannelPush:
-		s.ChannelSend(context.Background(), ev.Channel, ev.Data)
+		s.ChannelSend(ctx, ev.Channel, ev.Data)
 	}
 	return core.Result{OK: true}
 }