Add Heap.enable + pipeline health diagnostic to heap tracking

- Call Heap.enable before Heap.startTracking (required for events)
- Add pre-GC to reduce snapshot size before starting tracking
- Wait for trackingStart event to confirm pipeline health
- Report pipeline status (healthy/warning) from heap_start_tracking
- Update tests with mock handlers for Heap.enable and Heap.gc

Author: nnemirovsky

cmd/iwdp-mcp/main.go

@@ -1715,7 +1715,7 @@ func registerTools(server *mcp.Server) {
 	})
 
 	mcp.AddTool(server, &mcp.Tool{
-		Name: "heap_start_tracking", Description: "Start tracking heap allocations — collects heap snapshot events",
+		Name: "heap_start_tracking", Description: "Start tracking heap allocations — collects GC events. Waits up to 5s to confirm event pipeline health.",
 	}, func(ctx context.Context, req *mcp.CallToolRequest, _ EmptyInput) (*mcp.CallToolResult, any, error) {
 		c, err := getClient(ctx)
 		if err != nil {
@@ -1727,7 +1727,19 @@ func registerTools(server *mcp.Server) {
 		}
 		collector := sess.heapTrackingCollector
 		sess.mu.Unlock()
-		return nil, ok(), collector.Start(ctx, c)
+		if err := collector.Start(ctx, c); err != nil {
+			return nil, OKOutput{}, err
+		}
+		if collector.PipelineHealthy() {
+			return nil, struct {
+				OK      bool   `json:"ok"`
+				Message string `json:"message"`
+			}{true, "Heap tracking started. Event pipeline confirmed healthy — GC events will be captured."}, nil
+		}
+		return nil, struct {
+			OK      bool   `json:"ok"`
+			Warning string `json:"warning"`
+		}{true, "Heap tracking started, but trackingStart event not received (iwdp may not support the 50-200MB snapshot relay). GC events may not be captured. Use heap_snapshot and heap_gc directly instead."}, nil
 	})
 
 	mcp.AddTool(server, &mcp.Tool{

internal/tools/memory.go

@@ -186,30 +186,47 @@ type GarbageCollection struct {
 	EndTime   float64 `json:"endTime"`
 }
 
+// HeapTrackingReadyTimeout is how long Start waits for the trackingStart event
+// to confirm the event pipeline is healthy. Default 5s.
+var HeapTrackingReadyTimeout = 5 * time.Second
+
 // HeapTrackingResult holds the collected heap tracking data.
-// Note: Heap.trackingStart/trackingComplete carry 50-200MB+ snapshot payloads
-// that crash iwdp's WebSocket relay, so we intentionally skip them and only
-// collect lightweight garbageCollected events. Use the dedicated heap_snapshot
-// tool for snapshots (it uses Heap.snapshot which returns data in-band).
 type HeapTrackingResult struct {
-	GCEvents []GarbageCollection `json:"gcEvents,omitempty"`
+	GCEvents        []GarbageCollection `json:"gcEvents,omitempty"`
+	PipelineHealthy bool                `json:"pipelineHealthy"`
 }
 
 // HeapTrackingCollector collects Heap GC events during tracking.
-// Snapshot events (trackingStart/trackingComplete) are intentionally ignored
-// because their 50-200MB+ payloads crash iwdp's WebSocket relay.
+//
+// Heap.startTracking triggers a trackingStart event carrying a full heap
+// snapshot (50-200MB+). If iwdp can relay this massive event successfully,
+// the event pipeline is healthy and subsequent garbageCollected events will
+// arrive. If not, the pipeline is broken and no events will be captured.
+// Start() waits for the trackingStart event to diagnose this.
 type HeapTrackingCollector struct {
-	mu       sync.Mutex
-	gcEvents []GarbageCollection
-	started  bool
+	mu              sync.Mutex
+	gcEvents        []GarbageCollection
+	started         bool
+	ready           chan struct{}
+	pipelineHealthy bool
 }
 
 // NewHeapTrackingCollector creates a new HeapTrackingCollector.
 func NewHeapTrackingCollector() *HeapTrackingCollector {
 	return &HeapTrackingCollector{}
 }
 
+// PipelineHealthy reports whether the event pipeline survived the massive
+// trackingStart event. If false, garbageCollected events won't arrive.
+func (c *HeapTrackingCollector) PipelineHealthy() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.pipelineHealthy
+}
+
 // Start begins heap tracking, collecting garbageCollected events.
+// It waits up to HeapTrackingReadyTimeout for the trackingStart event to
+// confirm the event pipeline is healthy.
 func (c *HeapTrackingCollector) Start(ctx context.Context, client *webkit.Client) error {
 	c.mu.Lock()
 	if c.started {
@@ -218,8 +235,27 @@ func (c *HeapTrackingCollector) Start(ctx context.Context, client *webkit.Client
 	}
 	c.started = true
 	c.gcEvents = nil
+	c.ready = make(chan struct{})
+	c.pipelineHealthy = false
 	c.mu.Unlock()
 
+	// trackingStart handler: signals that the massive snapshot event arrived
+	// at our client, confirming the event pipeline is healthy. We discard
+	// the snapshot data (50-200MB+) — use heap_snapshot for snapshots.
+	client.OnEvent("Heap.trackingStart", func(method string, params json.RawMessage) {
+		c.mu.Lock()
+		c.pipelineHealthy = true
+		ch := c.ready
+		c.mu.Unlock()
+		if ch != nil {
+			select {
+			case <-ch:
+			default:
+				close(ch)
+			}
+		}
+	})
+
 	client.OnEvent("Heap.garbageCollected", func(method string, params json.RawMessage) {
 		var evt struct {
 			Collection GarbageCollection `json:"collection"`
@@ -231,8 +267,34 @@ func (c *HeapTrackingCollector) Start(ctx context.Context, client *webkit.Client
 		}
 	})
 
+	// Enable Heap domain — required for events to be dispatched.
+	_, _ = client.Send(ctx, "Heap.enable", nil)
+
+	// Pre-GC: reduce heap size before startTracking to minimize the snapshot
+	// payload in the trackingStart event. A smaller snapshot is more likely
+	// to survive iwdp's WebSocket relay (which has a 64MB message limit).
+	_, _ = client.Send(ctx, "Heap.gc", nil)
+
 	_, err := client.Send(ctx, "Heap.startTracking", nil)
-	return err
+	if err != nil {
+		return err
+	}
+
+	// Wait for trackingStart event to confirm the event pipeline survived
+	// the massive snapshot payload. If it arrives, GC events will work.
+	// If not, iwdp couldn't relay the 50-200MB+ event and the pipeline is broken.
+	c.mu.Lock()
+	ch := c.ready
+	c.mu.Unlock()
+	select {
+	case <-ch:
+		// Pipeline healthy — trackingStart event arrived.
+	case <-time.After(HeapTrackingReadyTimeout):
+		// Timeout — event pipeline likely broken by massive snapshot.
+	case <-ctx.Done():
+	}
+
+	return nil
 }
 
 // Stop stops heap tracking and returns collected GC events.
@@ -244,6 +306,7 @@ func (c *HeapTrackingCollector) Stop(ctx context.Context, client *webkit.Client)
 		c.mu.Unlock()
 		return &HeapTrackingResult{}, nil
 	}
+	healthy := c.pipelineHealthy
 	c.mu.Unlock()
 
 	// stopTracking triggers trackingComplete with a massive snapshot payload
@@ -253,7 +316,8 @@ func (c *HeapTrackingCollector) Stop(ctx context.Context, client *webkit.Client)
 	c.mu.Lock()
 	c.started = false
 	result := &HeapTrackingResult{
-		GCEvents: make([]GarbageCollection, len(c.gcEvents)),
+		GCEvents:        make([]GarbageCollection, len(c.gcEvents)),
+		PipelineHealthy: healthy,
 	}
 	copy(result.GCEvents, c.gcEvents)
 	c.gcEvents = nil

internal/tools/tools_domains_test.go

@@ -575,18 +575,33 @@ func TestHeapSnapshot(t *testing.T) {
 
 func TestHeapTrackingCollector(t *testing.T) {
 	mock, client := setup(t)
+	mock.HandleFunc("Heap.enable", map[string]interface{}{})
+	mock.HandleFunc("Heap.gc", map[string]interface{}{})
 	mock.HandleFunc("Heap.startTracking", map[string]interface{}{})
 	mock.HandleFunc("Heap.stopTracking", map[string]interface{}{})
 
 	collector := tools.NewHeapTrackingCollector()
 	ctx := context.Background()
+
+	// Send trackingStart shortly after startTracking to confirm pipeline health.
+	// In production, this event carries 50-200MB+ snapshot data; here we use a
+	// small mock payload. Start() waits for this signal before returning.
+	go func() {
+		time.Sleep(50 * time.Millisecond)
+		_ = mock.SendEvent("Heap.trackingStart", map[string]interface{}{
+			"timestamp":    1000.0,
+			"snapshotData": "mock-snapshot",
+		})
+	}()
+
 	if err := collector.Start(ctx, client); err != nil {
 		t.Fatalf("HeapTrackingCollector.Start: %v", err)
 	}
+	if !collector.PipelineHealthy() {
+		t.Error("expected pipeline to be healthy after trackingStart event")
+	}
 
-	// Send garbageCollected events (the only events we collect — snapshot
-	// events from trackingStart/trackingComplete are intentionally ignored
-	// because their 50-200MB payloads crash iwdp).
+	// Send garbageCollected events.
 	if err := mock.SendEvent("Heap.garbageCollected", map[string]interface{}{
 		"collection": map[string]interface{}{
 			"type":      "full",
@@ -611,6 +626,9 @@ func TestHeapTrackingCollector(t *testing.T) {
 	if err != nil {
 		t.Fatalf("HeapTrackingCollector.Stop: %v", err)
 	}
+	if !result.PipelineHealthy {
+		t.Error("expected PipelineHealthy=true in result")
+	}
 	if len(result.GCEvents) != 2 {
 		t.Fatalf("expected 2 GC events, got %d", len(result.GCEvents))
 	}